Profiler

External package

This page documents calibrax, the benchmarking library datarax depends on.

GPU memory profiling, hardware-adaptive optimization, and memory analysis for Datarax pipelines.

See Also

• Benchmarking Overview - All benchmarking tools
• Performance Tools - Optimization
• Benchmarking Guide

Overview

This module provides three components:

• GPUMemoryProfiler — Detects GPU availability and reports memory usage (used/total/utilization). Also analyzes memory patterns across multiple measurements to detect leaks and high utilization.
• MemoryOptimizer — Analyzes a pipeline function's memory footprint by measuring baseline, peak, and post-GC memory. Returns optimization suggestions.
• AdaptiveOperation — Auto-detects hardware (CPU/GPU/TPU) and configures optimal tile sizes, precision, and batch sizes. Also pads tensor shapes to hardware tile-size multiples via optimize_shapes().

Quick Start

Check GPU memory

from calibrax.profiling import GPUMemoryProfiler

profiler = GPUMemoryProfiler()
usage = profiler.get_memory_usage()
print(f"GPU memory: {usage['gpu_memory_used_mb']:.1f} / {usage['gpu_memory_total_mb']:.1f} MB")
print(f"Utilization: {usage.get('gpu_memory_utilization', 0):.1%}")

Analyze pipeline memory

from calibrax.profiling import MemoryOptimizer

optimizer = MemoryOptimizer()
analysis = optimizer.analyze_pipeline_memory(pipeline_fn, sample_data)
if analysis is not None:
    print(f"Peak usage: {analysis.peak_usage_mb:.1f} MB")
    print(f"Memory efficiency: {analysis.memory_efficiency:.1%}")
    for suggestion in analysis.suggestions:
        print(f"  - {suggestion}")

---

calibrax.profiling

Profiling: timing, resources, GPU, energy, FLOPs, hardware, roofline, compilation, complexity.

HARDWARE_SPECS module-attribute

HARDWARE_SPECS: dict[str, dict[str, Any]] = {'tpu_v5e': {'peak_flops': 197000000000000.0, 'peak_flops_bf16': 197000000000000.0, 'memory_bandwidth': 1600000000000.0, 'critical_intensity': 123.125}, 'a100_80g': {'peak_flops': 312000000000000.0, 'peak_flops_bf16': 312000000000000.0, 'memory_bandwidth': 2039000000000.0, 'critical_intensity': 153.0, 'tensor_core_shapes': [(16, 16, 16), (16, 16, 8)]}, 'h100': {'peak_flops': 989000000000000.0, 'peak_flops_bf16': 989000000000000.0, 'memory_bandwidth': 3350000000000.0, 'critical_intensity': 295.0, 'tensor_core_shapes': [(16, 16, 16)]}, 'cpu_generic': {'peak_flops': 2000000000000.0, 'peak_flops_bf16': 2000000000000.0, 'memory_bandwidth': 200000000000.0, 'critical_intensity': 10.0, 'simd_width': 8}}

__all__ module-attribute

__all__ = ['CompilationProfiler', 'CompilationResult', 'XLAOptimizationResult', 'ComplexityResult', 'analyze_complexity', 'EnergyMonitor', 'EnergySample', 'EnergySummary', 'FlopsCounter', 'FlopsResult', 'AdaptiveOperation', 'GPUMemoryProfiler', 'HardwareConfig', 'MemoryAnalysis', 'MemoryOptimizer', 'HARDWARE_SPECS', 'detect_hardware_specs', 'measure_execution_time', 'GPUProfilerProtocol', 'ResourceMonitor', 'ResourceSample', 'ResourceSummary', 'RooflineAnalyzer', 'RooflineResult', 'TimingCollector', 'TimingSample', 'TraceLinker', 'TraceReference']

CompilationProfiler

CompilationProfiler()

Analyzes JAX JIT compilation performance and optimization.

Instruments JIT-compiled functions to track compilation cache hits/misses, compilation times, and input shape consistency. Use profile_jit_compilation to wrap a function, then call get_result() for aggregated analysis.

_compilation_cache instance-attribute

_compilation_cache: dict[str, dict[str, Any]] = {}

_compilation_stats instance-attribute

_compilation_stats: dict[str, list[dict[str, Any]]] = defaultdict(list)

_cache_hit_count instance-attribute

_cache_hit_count = 0

_cache_miss_count instance-attribute

_cache_miss_count = 0

profile_jit_compilation

profile_jit_compilation(func: Callable[..., Any]) -> Callable[..., Any]

Create an instrumented wrapper that profiles JIT compilation.

The returned callable tracks cache hits/misses, compilation times, and input shape patterns. Results accumulate in this profiler instance.

Parameters:

Name	Type	Description	Default
func	Callable[..., Any]	JAX function to instrument.	required

Returns:

Type	Description
Callable[..., Any]	Instrumented function with identical signature.

get_result

get_result() -> CompilationResult

Get aggregated compilation profiling results.

Returns:

Type	Description
CompilationResult	CompilationResult with cache statistics, timing, and recommendations.

estimate_xla_optimization

estimate_xla_optimization(func: Callable[..., Any], *sample_args: Any) -> XLAOptimizationResult

Estimate XLA optimization effectiveness by analyzing HLO text.

Parameters:

Name	Type	Description	Default
func	Callable[..., Any]	JAX function to analyze.	required
*sample_args	Any	Example arguments for lowering/compiling.	()

Returns:

Type	Description
XLAOptimizationResult	XLAOptimizationResult with HLO analysis metrics.

reset

reset() -> None

Reset all profiling state.

_create_function_signature

_create_function_signature(func: Callable[..., Any], args: tuple[Any, ...], kwargs: dict[str, Any]) -> str

Create a unique signature for function + argument shapes.

Parameters:

Name	Type	Description	Default
func	Callable[..., Any]	The function being called.	required
args	tuple[Any, ...]	Positional arguments.	required
kwargs	dict[str, Any]	Keyword arguments.	required

Returns:

Type	Description
str	MD5 hex digest identifying the function + input signature.

_generate_recommendations

_generate_recommendations(cache_hit_rate: float, avg_compilation_time: float) -> list[str]

Generate compilation optimization recommendations.

Parameters:

Name	Type	Description	Default
cache_hit_rate	float	Fraction of cache hits.	required
avg_compilation_time	float	Average compilation time in seconds.	required

Returns:

Type	Description
list[str]	List of recommendation strings.

_assess_health

_assess_health(cache_hit_rate: float, avg_compilation_time: float) -> tuple[float, str]

Assess overall compilation health.

Parameters:

Name	Type	Description	Default
cache_hit_rate	float	Fraction of cache hits.	required
avg_compilation_time	float	Average compilation time in seconds.	required

Returns:

Type	Description
tuple[float, str]	Tuple of (health_score, health_level).

_analyze_hlo

_analyze_hlo(hlo_text: str) -> XLAOptimizationResult

Analyze HLO text for optimization patterns.

Parameters:

Name	Type	Description	Default
hlo_text	str	HLO text from compiled XLA module.	required

Returns:

Type	Description
XLAOptimizationResult	XLAOptimizationResult with analysis metrics.

_calculate_optimization_score

_calculate_optimization_score(fusion_ratio: float, arithmetic_ratio: float, memory_ratio: float) -> float

Calculate an optimization effectiveness score (0-1).

Parameters:

Name	Type	Description	Default
fusion_ratio	float	Fraction of fused kernels.	required
arithmetic_ratio	float	Fraction of arithmetic operations.	required
memory_ratio	float	Fraction of memory operations.	required

Returns:

Type	Description
float	Score between 0 and 1.

_generate_xla_recommendations

_generate_xla_recommendations(fusion_ratio: float, arithmetic_ratio: float, memory_ratio: float) -> list[str]

Generate XLA optimization recommendations.

Parameters:

Name	Type	Description	Default
fusion_ratio	float	Fraction of fused kernels.	required
arithmetic_ratio	float	Fraction of arithmetic operations.	required
memory_ratio	float	Fraction of memory operations.	required

Returns:

Type	Description
list[str]	List of recommendation strings.

CompilationResult dataclass

CompilationResult(*, cache_hit_rate: float, total_calls: int, cache_hits: int, cache_misses: int, avg_compilation_time_ms: float, max_compilation_time_ms: float, unique_signatures: int, health_score: float, health_level: str, recommendations: tuple[str, ...] = ())

Result of compilation profiling analysis.

Attributes:

Name	Type	Description
cache_hit_rate	float	Fraction of calls that hit the compilation cache.
total_calls	int	Total number of profiled function calls.
cache_hits	int	Number of cache hits.
cache_misses	int	Number of cache misses (triggering compilation).
avg_compilation_time_ms	float	Average compilation time in milliseconds.
max_compilation_time_ms	float	Maximum compilation time in milliseconds.
unique_signatures	int	Number of unique function signatures compiled.
health_score	float	Overall compilation health score (0-1).
health_level	str	Human-readable health level.
recommendations	tuple[str, ...]	Optimization recommendations.

cache_hit_rate instance-attribute

cache_hit_rate: float

total_calls instance-attribute

total_calls: int

cache_hits instance-attribute

cache_hits: int

cache_misses instance-attribute

cache_misses: int

avg_compilation_time_ms instance-attribute

avg_compilation_time_ms: float

max_compilation_time_ms instance-attribute

max_compilation_time_ms: float

unique_signatures instance-attribute

unique_signatures: int

health_score instance-attribute

health_score: float

health_level instance-attribute

health_level: str

recommendations class-attribute instance-attribute

recommendations: tuple[str, ...] = ()

to_dict

to_dict() -> dict[str, Any]

Serialize to a JSON-compatible dictionary.

from_dict classmethod

from_dict(data: dict[str, Any]) -> CompilationResult

Deserialize from a dictionary.

Parameters:

Name	Type	Description	Default
data	dict[str, Any]	Dictionary with compilation result fields.	required

Returns:

Type	Description
CompilationResult	Reconstructed CompilationResult instance.

XLAOptimizationResult dataclass

XLAOptimizationResult(*, optimization_score: float, fusion_ratio: float, arithmetic_ratio: float, memory_ratio: float, total_kernels: int, recommendations: tuple[str, ...] = ())

Result of XLA optimization effectiveness analysis.

Attributes:

Name	Type	Description
optimization_score	float	Overall optimization score (0-1).
fusion_ratio	float	Fraction of fused kernels.
arithmetic_ratio	float	Fraction of arithmetic operations.
memory_ratio	float	Fraction of memory operations.
total_kernels	int	Total number of HLO kernels.
recommendations	tuple[str, ...]	Optimization recommendations.

optimization_score instance-attribute

optimization_score: float

fusion_ratio instance-attribute

fusion_ratio: float

arithmetic_ratio instance-attribute

arithmetic_ratio: float

memory_ratio instance-attribute

memory_ratio: float

total_kernels instance-attribute

total_kernels: int

recommendations class-attribute instance-attribute

recommendations: tuple[str, ...] = ()

to_dict

to_dict() -> dict[str, Any]

Serialize to a JSON-compatible dictionary.

from_dict classmethod

from_dict(data: dict[str, Any]) -> XLAOptimizationResult

Deserialize from a dictionary.

Parameters:

Name	Type	Description	Default
data	dict[str, Any]	Dictionary with XLA optimization result fields.	required

Returns:

Type	Description
XLAOptimizationResult	Reconstructed XLAOptimizationResult instance.

ComplexityResult dataclass

ComplexityResult(*, total_parameters: int, parameter_memory_mb: float, largest_layer_name: str, largest_layer_params: int, input_shape: tuple[int, ...], estimated_memory_mb: float, total_estimated_operations: int, dominant_complexity: str, scaling_characteristics: dict[str, str] = dict())

Result of model complexity analysis.

Attributes:

Name	Type	Description
total_parameters	int	Total number of trainable parameters.
parameter_memory_mb	float	Memory consumed by parameters (float32).
largest_layer_name	str	Name of the layer with the most parameters.
largest_layer_params	int	Parameter count of the largest layer.
input_shape	tuple[int, ...]	Shape of the analyzed input.
estimated_memory_mb	float	Estimated total memory (params + activations).
total_estimated_operations	int	Estimated total operations count.
dominant_complexity	str	Name of the dominant operation type.
scaling_characteristics	dict[str, str]	Mapping of operation type to complexity class.

total_parameters instance-attribute

total_parameters: int

parameter_memory_mb instance-attribute

parameter_memory_mb: float

largest_layer_name instance-attribute

largest_layer_name: str

largest_layer_params instance-attribute

largest_layer_params: int

input_shape instance-attribute

input_shape: tuple[int, ...]

estimated_memory_mb instance-attribute

estimated_memory_mb: float

total_estimated_operations instance-attribute

total_estimated_operations: int

dominant_complexity instance-attribute

dominant_complexity: str

scaling_characteristics class-attribute instance-attribute

scaling_characteristics: dict[str, str] = field(default_factory=dict)

to_dict

to_dict() -> dict[str, Any]

Serialize to a JSON-compatible dictionary.

from_dict classmethod

from_dict(data: dict[str, Any]) -> ComplexityResult

Deserialize from a dictionary.

Parameters:

Name	Type	Description	Default
data	dict[str, Any]	Dictionary with complexity result fields.	required

Returns:

Type	Description
ComplexityResult	Reconstructed ComplexityResult instance.

EnergyMonitor

EnergyMonitor(sample_interval_sec: float = 0.1)

Background energy monitoring via NVML and RAPL.

Uses daemon thread sampling at configurable interval. Gracefully degrades when NVML or RAPL is unavailable.

Usage:

with EnergyMonitor() as mon:
    # ... run benchmark ...
summary = mon.summary

Parameters:

Name	Type	Description	Default
sample_interval_sec	float	Seconds between energy samples.	0.1

_interval instance-attribute

_interval = sample_interval_sec

_samples instance-attribute

_samples: list[EnergySample] = []

_sampling_thread instance-attribute

_sampling_thread = SamplingThread(target=self._sample_loop)

_gpu_power_readings instance-attribute

_gpu_power_readings: list[tuple[float, float]] = []

_rapl_start_uj instance-attribute

_rapl_start_uj: int | None = None

_rapl_prev_uj instance-attribute

_rapl_prev_uj: int | None = None

_rapl_accumulated_uj instance-attribute

_rapl_accumulated_uj: int = 0

samples property

samples: list[EnergySample]

Return a copy of all collected samples.

summary property

summary: EnergySummary

Compute aggregated energy summary.

Returns:

Type	Description
EnergySummary	EnergySummary with totals, or None fields when unavailable.

__enter__

__enter__() -> EnergyMonitor

Start background energy sampling thread.

__exit__

__exit__(*args: object) -> None

Stop background energy sampling thread.

_sample_loop

_sample_loop() -> None

Collect energy samples until stopped.

_read_gpu_power

_read_gpu_power() -> float | None

Read GPU power in watts.

_update_rapl

_update_rapl() -> None

Update RAPL accumulated energy with wraparound handling.

_compute_gpu_energy

_compute_gpu_energy() -> float | None

Compute cumulative GPU energy via rectangular integration.

_compute_cpu_energy

_compute_cpu_energy() -> float | None

Compute CPU energy from RAPL counters.

_compute_gpu_power_stats

_compute_gpu_power_stats() -> tuple[float | None, float | None]

Compute mean and peak GPU power from samples.

Returns:

Type	Description
tuple[float | None, float | None]	Tuple of (mean_power, peak_power), both None if no GPU data.

EnergySample dataclass

EnergySample(*, timestamp: float, gpu_power_watts: float | None, cpu_energy_joules: float | None, gpu_energy_joules: float | None)

Single energy measurement at a point in time.

Attributes:

Name	Type	Description
timestamp	float	Time of measurement (perf_counter).
gpu_power_watts	float | None	Instantaneous GPU power (None if unavailable).
cpu_energy_joules	float | None	Cumulative CPU energy since monitoring start.
gpu_energy_joules	float | None	Cumulative GPU energy since monitoring start.

timestamp instance-attribute

timestamp: float

gpu_power_watts instance-attribute

gpu_power_watts: float | None

cpu_energy_joules instance-attribute

cpu_energy_joules: float | None

gpu_energy_joules instance-attribute

gpu_energy_joules: float | None

EnergySummary dataclass

EnergySummary(*, total_gpu_energy_joules: float | None, total_cpu_energy_joules: float | None, total_combined_energy_joules: float | None, mean_gpu_power_watts: float | None, peak_gpu_power_watts: float | None, duration_sec: float, num_samples: int)

Aggregated energy usage over a monitoring period.

Attributes:

Name	Type	Description
total_gpu_energy_joules	float | None	Total GPU energy consumed.
total_cpu_energy_joules	float | None	Total CPU energy consumed.
total_combined_energy_joules	float | None	GPU + CPU combined.
mean_gpu_power_watts	float | None	Average GPU power draw.
peak_gpu_power_watts	float | None	Maximum GPU power draw.
duration_sec	float	Monitoring duration.
num_samples	int	Total samples collected.

total_gpu_energy_joules instance-attribute

total_gpu_energy_joules: float | None

total_cpu_energy_joules instance-attribute

total_cpu_energy_joules: float | None

total_combined_energy_joules instance-attribute

total_combined_energy_joules: float | None

mean_gpu_power_watts instance-attribute

mean_gpu_power_watts: float | None

peak_gpu_power_watts instance-attribute

peak_gpu_power_watts: float | None

duration_sec instance-attribute

duration_sec: float

num_samples instance-attribute

num_samples: int

FlopsCounter

Count FLOPs of JAX functions via jaxpr analysis.

Uses jax.make_jaxpr to trace the function and counts FLOPs for each primitive based on operation-specific rules.

For NNX models that use stochastic operations (dropout, etc.), use flax.nnx.tabulate(model, *args, compute_flops=True) instead — it handles NNX state management internally.

count

count(fn: Callable[..., Any], *args: Any, static_argnums: tuple[int, ...] = ()) -> FlopsResult

Count FLOPs for a function with given example arguments.

Parameters:

Name	Type	Description	Default
fn	Callable[..., Any]	JAX function to analyze.	required
*args	Any	Example arguments for tracing.	()
static_argnums	tuple[int, ...]	Argument indices to treat as static.	()

Returns:

Type	Description
FlopsResult	FlopsResult with FLOP count and breakdown.

_count_jaxpr

_count_jaxpr(jaxpr: Jaxpr, flops_by_op: dict[str, int]) -> int

Recursively count FLOPs in a Jaxpr.

Parameters:

Name	Type	Description	Default
jaxpr	Jaxpr	The Jaxpr to analyze.	required
flops_by_op	dict[str, int]	Accumulator for per-operation FLOP counts.	required

Returns:

Type	Description
int	Total FLOPs in this Jaxpr.

_count_eqn

_count_eqn(eqn: JaxprEqn, flops_by_op: dict[str, int]) -> int

Count FLOPs for a single Jaxpr equation.

Parameters:

Name	Type	Description	Default
eqn	JaxprEqn	The equation to analyze.	required
flops_by_op	dict[str, int]	Accumulator for per-operation FLOP counts.	required

Returns:

Type	Description
int	FLOPs for this equation.

_classify_primitive_flops

_classify_primitive_flops(name: str, eqn: JaxprEqn, flops_by_op: dict[str, int]) -> int

Dispatch primitive to appropriate FLOP counting strategy.

Parameters:

Name	Type	Description	Default
name	str	Primitive operation name.	required
eqn	JaxprEqn	The Jaxpr equation.	required
flops_by_op	dict[str, int]	Accumulator for nested operations.	required

Returns:

Type	Description
int	Estimated FLOPs for this primitive.

_count_dot_general

_count_dot_general(eqn: JaxprEqn) -> int

Count FLOPs for dot_general (matmul).

For (M, K) @ (K, N) -> 2 * M * K * N.

_count_conv

_count_conv(eqn: JaxprEqn) -> int

Count FLOPs for conv_general_dilated.

_count_nested_jaxpr

_count_nested_jaxpr(eqn: JaxprEqn, flops_by_op: dict[str, int]) -> int

Count FLOPs in nested Jaxprs (pjit, scan, etc.).

_output_size

_output_size(eqn: JaxprEqn) -> int

Product of output shape dimensions.

_input_size

_input_size(eqn: JaxprEqn) -> int

Product of first input shape dimensions.

FlopsResult dataclass

FlopsResult(*, total_flops: int, flops_by_operation: dict[str, int], num_operations: int, function_name: str)

Result of FLOP counting for a function.

Attributes:

Name	Type	Description
total_flops	int	Total estimated FLOPs.
flops_by_operation	dict[str, int]	Breakdown by primitive operation name.
num_operations	int	Number of JAX primitives in the trace.
function_name	str	Name of the analyzed function.

total_flops instance-attribute

total_flops: int

flops_by_operation instance-attribute

flops_by_operation: dict[str, int]

num_operations instance-attribute

num_operations: int

function_name instance-attribute

function_name: str

AdaptiveOperation

AdaptiveOperation()

Hardware-adaptive operations with auto-detection.

Detects the current JAX backend (CPU/GPU/TPU) and provides optimized configuration and shape padding.

config instance-attribute

config = self._detect_hardware()

_detect_hardware

_detect_hardware() -> HardwareConfig

Detect hardware and return optimal configuration.

Returns:

Type	Description
HardwareConfig	HardwareConfig for the detected platform.

_detect_gpu_config

_detect_gpu_config() -> HardwareConfig

Detect GPU variant and return config.

Returns:

Type	Description
HardwareConfig	HardwareConfig for the detected GPU, or CPU default on failure.

optimize_shapes

optimize_shapes(*shapes: tuple[int, ...]) -> list[tuple[int, ...]]

Pad tensor shapes to align with hardware tile size.

Parameters:

Name	Type	Description	Default
*shapes	tuple[int, ...]	Variable number of tensor shapes to optimize.	()

Returns:

Type	Description
list[tuple[int, ...]]	List of optimized shapes padded to tile_size multiples.

GPUMemoryProfiler

GPUMemoryProfiler()

GPU memory profiling satisfying GPUProfilerProtocol.

Uses multi-fallback strategy: memory_stats -> xla_bridge -> zeros.

has_gpu instance-attribute

has_gpu = len(jax.devices('gpu')) > 0

get_memory_usage

get_memory_usage() -> dict[str, float]

Get current GPU memory usage statistics.

Returns:

Type	Description
dict[str, float]	Dictionary with gpu_memory_used_mb, gpu_memory_total_mb,
dict[str, float]	and optionally gpu_memory_utilization.

get_utilization

get_utilization() -> float

Get GPU utilization percentage for ResourceMonitor.

Returns:

Type	Description
float	GPU memory utilization as percentage (0-100), or 0.0.

_safe_nvml_query

_safe_nvml_query(query_fn: Callable[[Any], dict[str, float]], fallback: dict[str, float]) -> dict[str, float]

Execute an NVML query with init and fallback on failure.

Parameters:

Name	Type	Description	Default
query_fn	Callable[[Any], dict[str, float]]	Function that takes an NVML handle and returns metrics.	required
fallback	dict[str, float]	Default dict to return on failure.	required

Returns:

Type	Description
dict[str, float]	Query result or fallback on any error.

get_clock_info

get_clock_info() -> dict[str, float]

Get current GPU clock frequencies via NVML.

Returns:

Type	Description
dict[str, float]	Dictionary with 'gpu_clock_mhz' and 'mem_clock_mhz' keys.
dict[str, float]	Returns zeros if NVML is unavailable or query fails.

get_power_info

get_power_info() -> dict[str, float]

Get current GPU power draw and limit via NVML.

Returns:

Type	Description
dict[str, float]	Dictionary with 'power_draw_w' and 'power_limit_w' keys.
dict[str, float]	Returns zeros if NVML is unavailable or query fails.

analyze_memory_pattern

analyze_memory_pattern(measurements: list[dict[str, float]]) -> list[str]

Analyze memory usage patterns and suggest optimizations.

Parameters:

Name	Type	Description	Default
measurements	list[dict[str, float]]	List of memory usage dictionaries.	required

Returns:

Type	Description
list[str]	List of optimization suggestion strings.

HardwareConfig dataclass

HardwareConfig(*, platform: str, precision: str, tile_size: int, critical_batch_size: int, memory_layout: str, use_vmem_optimization: bool)

Hardware-specific optimization configuration.

Attributes:

Name	Type	Description
platform	str	Detected platform ("cpu", "tpu", "gpu_modern", "gpu_legacy").
precision	str	Recommended floating-point precision string.
tile_size	int	Tile size for matrix operation alignment.
critical_batch_size	int	Optimal batch size for the platform.
memory_layout	str	Memory layout preference.
use_vmem_optimization	bool	Whether VMEM optimization is available.

platform instance-attribute

platform: str

precision instance-attribute

precision: str

tile_size instance-attribute

tile_size: int

critical_batch_size instance-attribute

critical_batch_size: int

memory_layout instance-attribute

memory_layout: str

use_vmem_optimization instance-attribute

use_vmem_optimization: bool

MemoryAnalysis dataclass

MemoryAnalysis(*, baseline_memory_mb: float, peak_memory_mb: float, peak_usage_mb: float, retained_memory_mb: float, memory_efficiency: float, suggestions: tuple[str, ...] = ())

Result of pipeline memory analysis.

Attributes:

Name	Type	Description
baseline_memory_mb	float	Memory usage before pipeline execution.
peak_memory_mb	float	Memory usage at peak during execution.
peak_usage_mb	float	Peak usage above baseline.
retained_memory_mb	float	Memory retained after GC.
memory_efficiency	float	Ratio of freed memory to peak usage.
suggestions	tuple[str, ...]	Optimization suggestions.

baseline_memory_mb instance-attribute

baseline_memory_mb: float

peak_memory_mb instance-attribute

peak_memory_mb: float

peak_usage_mb instance-attribute

peak_usage_mb: float

retained_memory_mb instance-attribute

retained_memory_mb: float

memory_efficiency instance-attribute

memory_efficiency: float

suggestions class-attribute instance-attribute

suggestions: tuple[str, ...] = ()

MemoryOptimizer

Memory optimization analysis for pipeline functions.

analyze_pipeline_memory

analyze_pipeline_memory(pipeline_fn: Callable[[Any], Any], sample_data: Any) -> MemoryAnalysis | None

Analyze memory usage of a pipeline function.

Parameters:

Name	Type	Description	Default
pipeline_fn	Callable[[Any], Any]	Function to analyze.	required
sample_data	Any	Sample input data.	required

Returns:

Type	Description
MemoryAnalysis | None	MemoryAnalysis with measurements and suggestions,
MemoryAnalysis | None	or None if the pipeline raises an exception.

_get_rss_mb

_get_rss_mb() -> float

Get current process RSS in MB.

_generate_suggestions

_generate_suggestions(peak_usage: float, retained_memory: float) -> list[str]

Generate memory optimization suggestions.

Parameters:

Name	Type	Description	Default
peak_usage	float	Peak memory usage in MB.	required
retained_memory	float	Retained memory after GC in MB.	required

Returns:

Type	Description
list[str]	List of suggestion strings.

RooflineAnalyzer dataclass

RooflineAnalyzer(*, hardware_specs: dict[str, Any] = detect_hardware_specs())

Analyzes operation performance against hardware roofline limits.

Uses measured execution time and estimated FLOPs to determine whether an operation is compute-bound or memory-bound, and how efficiently it uses the available hardware resources.

Attributes:

Name	Type	Description
hardware_specs	dict[str, Any]	Hardware specification dictionary (auto-detected if not provided).

hardware_specs class-attribute instance-attribute

hardware_specs: dict[str, Any] = field(default_factory=detect_hardware_specs)

analyze_operation

analyze_operation(func: Callable[..., Any], inputs: list[Array], *, flops_override: int | None = None) -> RooflineResult

Perform roofline analysis on a JAX operation.

Parameters:

Name	Type	Description	Default
func	Callable[..., Any]	JAX function to analyze.	required
inputs	list[Array]	Input arrays for the function.	required
flops_override	int | None	If provided, use this FLOP count instead of estimating. For accurate results, pass the output of FlopsCounter.count().	None

Returns:

Type	Description
RooflineResult	RooflineResult with bottleneck classification and recommendations.

_estimate_flops

_estimate_flops(func: Callable[..., Any], inputs: list[Array]) -> int

Estimate FLOPs using XLA cost analysis when possible.

Falls back to a simple heuristic if cost_analysis is unavailable.

Parameters:

Name	Type	Description	Default
func	Callable[..., Any]	JAX function.	required
inputs	list[Array]	Input arrays.	required

Returns:

Type	Description
int	Estimated FLOP count.

_estimate_memory_traffic

_estimate_memory_traffic(func: Callable[..., Any], inputs: list[Array]) -> int

Estimate total memory traffic (input + output bytes).

Parameters:

Name	Type	Description	Default
func	Callable[..., Any]	JAX function.	required
inputs	list[Array]	Input arrays.	required

Returns:

Type	Description
int	Estimated bytes of memory traffic.

_generate_recommendations

_generate_recommendations(arithmetic_intensity: float, efficiency: float, bottleneck: str, inputs: list[Array], achieved_flops: float) -> list[str]

Generate optimization recommendations based on roofline analysis.

Parameters:

Name	Type	Description	Default
arithmetic_intensity	float	Achieved FLOPs per byte.	required
efficiency	float	Utilization of the binding resource.	required
bottleneck	str	"memory_bandwidth" or "compute".	required
inputs	list[Array]	Input arrays.	required
achieved_flops	float	Achieved FLOP/s.	required

Returns:

Type	Description
list[str]	List of recommendation strings.

RooflineResult dataclass

RooflineResult(*, arithmetic_intensity: float, critical_intensity: float, memory_bandwidth_utilization: float, flops_utilization: float, bottleneck: str, efficiency: float, execution_time_ms: float, recommendations: tuple[str, ...] = ())

Result of a roofline analysis on a JAX operation.

Attributes:

Name	Type	Description
arithmetic_intensity	float	Achieved FLOPs per byte of memory traffic.
critical_intensity	float	Hardware's ridge point (FLOPs/byte).
memory_bandwidth_utilization	float	Fraction of peak memory bandwidth used.
flops_utilization	float	Fraction of peak FLOPs achieved.
bottleneck	str	Either "memory_bandwidth" or "compute".
efficiency	float	Utilization of the binding resource.
execution_time_ms	float	Measured execution time in milliseconds.
recommendations	tuple[str, ...]	Optimization suggestions.

arithmetic_intensity instance-attribute

arithmetic_intensity: float

critical_intensity instance-attribute

critical_intensity: float

memory_bandwidth_utilization instance-attribute

memory_bandwidth_utilization: float

flops_utilization instance-attribute

flops_utilization: float

bottleneck instance-attribute

bottleneck: str

efficiency instance-attribute

efficiency: float

execution_time_ms instance-attribute

execution_time_ms: float

recommendations class-attribute instance-attribute

recommendations: tuple[str, ...] = ()

to_dict

to_dict() -> dict[str, Any]

Serialize to a JSON-compatible dictionary.

from_dict classmethod

from_dict(data: dict[str, Any]) -> RooflineResult

Deserialize from a dictionary.

Parameters:

Name	Type	Description	Default
data	dict[str, Any]	Dictionary with roofline result fields.	required

Returns:

Type	Description
RooflineResult	Reconstructed RooflineResult instance.

TimingCollector

TimingCollector(sync_fn: Callable[[Any], object] | None = None, warmup_iterations: int = 0)

Framework-agnostic timing with configurable GPU sync support.

Uses time.perf_counter() exclusively for accurate benchmarking. Supports configurable result synchronization via sync_fn and warm-up iteration exclusion for JIT-compiled workloads.

JAX dispatches operations asynchronously -- the host returns immediately while the device is still computing. Without an explicit synchronization barrier, perf_counter measures only host-side dispatch latency, not actual compute time. Pass a sync_fn that calls block_until_ready() on the workload result to force the host to wait for device completion before recording the timestamp.

Example -- JAX GPU timing with warm-up:

import jax.numpy as jnp

def run_step(batch):
    return jax.jit(step_fn)(batch)

collector = TimingCollector(
    sync_fn=lambda result: result.block_until_ready(),
    warmup_iterations=2,
)
sample = collector.measure_iteration(data_iter, num_batches=50, process_fn=run_step)
# sample.per_batch_times excludes the first 2 batches

Parameters:

Name	Type	Description	Default
sync_fn	Callable[[Any], object] | None	Synchronization function called with each batch result. For JAX: lambda result: result.block_until_ready() For PyTorch: lambda _: torch.cuda.synchronize() For CPU-only: None (default, no-op)	None
warmup_iterations	int	Number of initial batches to exclude from per_batch_times statistics. They are still executed (important for JIT warm-up) but omitted from the timing result. Default: 0.	0

Parameters:

Name	Type	Description	Default
sync_fn	Callable[[Any], object] | None	Synchronization function called with each batch result.	None
warmup_iterations	int	Number of initial batches to exclude from timing stats.	0

_sync_fn instance-attribute

_sync_fn = sync_fn or (lambda _result: None)

_warmup_iterations instance-attribute

_warmup_iterations = warmup_iterations

measure_iteration

measure_iteration(iterator: Iterator[Any], num_batches: int | None = None, process_fn: Callable[[Any], Any] | None = None, count_fn: Callable[[Any], int] | None = None) -> TimingSample

Measure timing for batches from an iterator.

Warm-up batches (if configured) are executed but excluded from per_batch_times. wall_clock_sec covers the entire run including warm-up. num_batches reflects total batches consumed.

Parameters:

Name	Type	Description	Default
iterator	Iterator[Any]	Data iterator to measure.	required
num_batches	int | None	Max batches to consume (including warmup). None exhausts iterator.	None
process_fn	Callable[[Any], Any] | None	Optional per-batch function whose execution is timed. Defaults to identity (the yielded batch is treated as result).	None
count_fn	Callable[[Any], int] | None	Function to count elements per batch. Default: 1 per batch.	None

Returns:

Type	Description
TimingSample	TimingSample with timing measurements.

measure_compilation_time

measure_compilation_time(fn: Callable[..., Any], *args: Any) -> float

Measure JIT compilation time for a JAX function.

Calls jax.jit(fn).lower(*args).compile() and times it. This measures the XLA compilation step only, not execution.

Parameters:

Name	Type	Description	Default
fn	Callable[..., Any]	JAX function to compile.	required
*args	Any	Example arguments for lowering.	()

Returns:

Type	Description
float	Compilation time in seconds.

TraceLinker

Links JAX profiler traces to benchmark runs.

Usage:

linker = TraceLinker()
with linker.trace("/tmp/my_trace") as ref:
    # ... run workload ...
print(ref.trace_dir)  # "/tmp/my_trace"

trace

trace(log_dir: str | Path, *, run_id: str | None = None, create_perfetto_link: bool = False, create_perfetto_trace: bool = False) -> Any

Start an XLA profiling session and record output metadata.

Wraps jax.profiler.trace() and records the output directory path as a TraceReference for downstream Store linkage.

Parameters:

Name	Type	Description	Default
log_dir	str | Path	Directory for trace output files.	required
run_id	str | None	Optional benchmark run ID to associate with the trace.	None
create_perfetto_link	bool	Whether to create a Perfetto link (passed to JAX).	False
create_perfetto_trace	bool	Whether to create a Perfetto trace (passed to JAX).	False

Yields:

Type	Description
Any	TraceReference with the trace directory and optional run ID.

TraceReference dataclass

TraceReference(*, trace_dir: str, run_id: str | None = None)

Reference to a JAX profiler trace output.

Attributes:

Name	Type	Description
trace_dir	str	Directory where the trace files were written.
run_id	str | None	Optional benchmark run ID to link the trace to.

trace_dir instance-attribute

trace_dir: str

run_id class-attribute instance-attribute

run_id: str | None = None

to_dict

to_dict() -> dict[str, Any]

Serialize to a JSON-compatible dictionary.

from_dict classmethod

from_dict(data: dict[str, Any]) -> TraceReference

Deserialize from a dictionary.

Parameters:

Name	Type	Description	Default
data	dict[str, Any]	Dictionary with trace reference fields.	required

Returns:

Type	Description
TraceReference	Reconstructed TraceReference instance.

analyze_complexity

analyze_complexity(model: Module, input_shape: tuple[int, ...]) -> ComplexityResult

Analyze complexity of a Flax NNX module.

Performs parameter counting, memory estimation, computational complexity analysis, and scaling characterization.

Parameters:

Name	Type	Description	Default
model	Module	Flax NNX model to analyze.	required
input_shape	tuple[int, ...]	Shape of input data (including batch dimension).	required

Returns:

Type	Description
ComplexityResult	ComplexityResult with detailed complexity metrics.

detect_hardware_specs

detect_hardware_specs() -> dict[str, Any]

Detect current hardware and return appropriate specifications.

Uses jax.default_backend() to determine the accelerator type and returns pre-configured specs for that platform.

Returns:

Type	Description
dict[str, Any]	Hardware specification dictionary with peak_flops, memory_bandwidth,
dict[str, Any]	and critical_intensity keys (among others).

measure_execution_time

measure_execution_time(func: Callable[..., Any], inputs: list[Array], warmup: int = 3, iterations: int = 10) -> float

Measure execution time of a JAX function with synchronization.

JIT-compiles the function, runs warmup iterations, then times iterations executions with block_until_ready() barriers.

Parameters:

Name	Type	Description	Default
func	Callable[..., Any]	JAX function to benchmark.	required
inputs	list[Array]	Input arguments as a list of arrays.	required
warmup	int	Number of warmup iterations (for JIT compilation).	3
iterations	int	Number of timed iterations.	10

Returns:

Type	Description
float	Average execution time in seconds.

_sampling

Shared sampling thread lifecycle for background monitors.

Provides a context-manager helper that encapsulates the daemon thread start/stop pattern used by ResourceMonitor and EnergyMonitor.

SamplingThread

SamplingThread(target: Callable[[], None])

Reusable daemon thread lifecycle for background sampling.

Usage:

thread = SamplingThread(target=self._sample_loop)
thread.start()   # in __enter__
thread.stop()    # in __exit__

Parameters:

Name	Type	Description	Default
target	Callable[[], None]	The sampling loop callable (runs in the daemon thread).	required

Parameters:

Name	Type	Description	Default
target	Callable[[], None]	Callable to run in the background thread.	required

_target instance-attribute

_target = target

_thread instance-attribute

_thread: Thread | None = None

stop_event instance-attribute

stop_event = threading.Event()

start

start() -> None

Clear stop event and start the daemon thread.

stop

stop() -> None

Signal the thread to stop and wait for it to finish.

carbon

Carbon emissions tracking via CodeCarbon integration.

Wraps the codecarbon.EmissionsTracker as a context manager, exposing emissions data as a frozen CarbonResult dataclass. Requires the optional codecarbon dependency (uv pip install "calibrax[codecarbon]").

CODECARBON_AVAILABLE module-attribute

CODECARBON_AVAILABLE = True

logger module-attribute

logger = logging.getLogger(__name__)

CarbonResult dataclass

CarbonResult(*, emissions_kg_co2: float, energy_consumed_kwh: float, duration_sec: float, country_iso_code: str | None = None)

Result of carbon emissions measurement.

Attributes:

Name	Type	Description
emissions_kg_co2	float	Total CO2 emissions in kilograms.
energy_consumed_kwh	float	Total energy consumed in kilowatt-hours.
duration_sec	float	Duration of the tracked period in seconds.
country_iso_code	str | None	ISO code of the country used for carbon intensity.

emissions_kg_co2 instance-attribute

emissions_kg_co2: float

energy_consumed_kwh instance-attribute

energy_consumed_kwh: float

duration_sec instance-attribute

duration_sec: float

country_iso_code class-attribute instance-attribute

country_iso_code: str | None = None

to_dict

to_dict() -> dict[str, Any]

Serialize to a JSON-compatible dictionary.

from_dict classmethod

from_dict(data: dict[str, Any]) -> CarbonResult

Deserialize from a dictionary.

Parameters:

Name	Type	Description	Default
data	dict[str, Any]	Dictionary with carbon result fields.	required

Returns:

Type	Description
CarbonResult	Reconstructed CarbonResult instance.

CarbonTracker

CarbonTracker(country_iso_code: str | None = None, log_level: str = 'warning')

Context manager for tracking carbon emissions via CodeCarbon.

Requires the codecarbon package. Install with:

uv pip install "calibrax[codecarbon]"

Usage:

with CarbonTracker() as tracker:
    # ... run workload ...
result = tracker.result()
print(f"Emissions: {result.emissions_kg_co2:.6f} kg CO2")

Parameters:

Name	Type	Description	Default
country_iso_code	str | None	Optional ISO country code for regional carbon intensity.	None
log_level	str	Logging level for CodeCarbon (default: "warning").	'warning'

Raises:

Type	Description
ImportError	If codecarbon is not installed.

Parameters:

Name	Type	Description	Default
country_iso_code	str | None	Optional ISO country code.	None
log_level	str	CodeCarbon logging level.	'warning'

Raises:

Type	Description
ImportError	If codecarbon is not installed.

_country_iso_code instance-attribute

_country_iso_code = country_iso_code

_log_level instance-attribute

_log_level = log_level

_tracker instance-attribute

_tracker: Any = None

_emissions instance-attribute

_emissions: float = 0.0

_energy instance-attribute

_energy: float = 0.0

_duration instance-attribute

_duration: float = 0.0

__enter__

__enter__() -> CarbonTracker

Start emissions tracking.

__exit__

__exit__(*args: object) -> None

Stop emissions tracking and record results.

result

result() -> CarbonResult

Get the carbon emissions result.

Call this after exiting the context manager.

Returns:

Type	Description
CarbonResult	CarbonResult with emissions, energy, and duration data.

compilation

JIT compilation profiler for JAX.

Analyzes JIT compilation efficiency, cache hit rates, XLA optimization effectiveness, and provides recommendations for compilation optimization.

logger module-attribute

logger = logging.getLogger(__name__)

_HLO_OPCODE_RE module-attribute

_HLO_OPCODE_RE = re.compile('\\b([A-Za-z_][\\w.\\-]*)\\s*\\(')

_FUSION_KIND_PATTERNS module-attribute

_FUSION_KIND_PATTERNS = ('kloop', 'kinput', 'koutput')

_MEMORY_OPS module-attribute

_MEMORY_OPS = ('copy', 'transpose', 'reshape', 'broadcast')

_ARITHMETIC_OPS module-attribute

_ARITHMETIC_OPS = ('add', 'multiply', 'dot', 'convolution', 'reduce')

_JAX_RUNTIME_ERROR module-attribute

_JAX_RUNTIME_ERROR: type[BaseException] = jax.errors.JaxRuntimeError

_RECOVERABLE_COMPILATION_ERRORS module-attribute

_RECOVERABLE_COMPILATION_ERRORS = (_JAX_RUNTIME_ERROR, AttributeError, OSError, RuntimeError, TypeError, ValueError)

CompilationResult dataclass

CompilationResult(*, cache_hit_rate: float, total_calls: int, cache_hits: int, cache_misses: int, avg_compilation_time_ms: float, max_compilation_time_ms: float, unique_signatures: int, health_score: float, health_level: str, recommendations: tuple[str, ...] = ())

Result of compilation profiling analysis.

Attributes:

Name	Type	Description
cache_hit_rate	float	Fraction of calls that hit the compilation cache.
total_calls	int	Total number of profiled function calls.
cache_hits	int	Number of cache hits.
cache_misses	int	Number of cache misses (triggering compilation).
avg_compilation_time_ms	float	Average compilation time in milliseconds.
max_compilation_time_ms	float	Maximum compilation time in milliseconds.
unique_signatures	int	Number of unique function signatures compiled.
health_score	float	Overall compilation health score (0-1).
health_level	str	Human-readable health level.
recommendations	tuple[str, ...]	Optimization recommendations.

cache_hit_rate instance-attribute

cache_hit_rate: float

total_calls instance-attribute

total_calls: int

cache_hits instance-attribute

cache_hits: int

cache_misses instance-attribute

cache_misses: int

avg_compilation_time_ms instance-attribute

avg_compilation_time_ms: float

max_compilation_time_ms instance-attribute

max_compilation_time_ms: float

unique_signatures instance-attribute

unique_signatures: int

health_score instance-attribute

health_score: float

health_level instance-attribute

health_level: str

recommendations class-attribute instance-attribute

recommendations: tuple[str, ...] = ()

to_dict

to_dict() -> dict[str, Any]

Serialize to a JSON-compatible dictionary.

from_dict classmethod

from_dict(data: dict[str, Any]) -> CompilationResult

Deserialize from a dictionary.

Parameters:

Name	Type	Description	Default
data	dict[str, Any]	Dictionary with compilation result fields.	required

Returns:

Type	Description
CompilationResult	Reconstructed CompilationResult instance.

XLAOptimizationResult dataclass

XLAOptimizationResult(*, optimization_score: float, fusion_ratio: float, arithmetic_ratio: float, memory_ratio: float, total_kernels: int, recommendations: tuple[str, ...] = ())

Result of XLA optimization effectiveness analysis.

Attributes:

Name	Type	Description
optimization_score	float	Overall optimization score (0-1).
fusion_ratio	float	Fraction of fused kernels.
arithmetic_ratio	float	Fraction of arithmetic operations.
memory_ratio	float	Fraction of memory operations.
total_kernels	int	Total number of HLO kernels.
recommendations	tuple[str, ...]	Optimization recommendations.

optimization_score instance-attribute

optimization_score: float

fusion_ratio instance-attribute

fusion_ratio: float

arithmetic_ratio instance-attribute

arithmetic_ratio: float

memory_ratio instance-attribute

memory_ratio: float

total_kernels instance-attribute

total_kernels: int

recommendations class-attribute instance-attribute

recommendations: tuple[str, ...] = ()

to_dict

to_dict() -> dict[str, Any]

Serialize to a JSON-compatible dictionary.

from_dict classmethod

from_dict(data: dict[str, Any]) -> XLAOptimizationResult

Deserialize from a dictionary.

Parameters:

Name	Type	Description	Default
data	dict[str, Any]	Dictionary with XLA optimization result fields.	required

Returns:

Type	Description
XLAOptimizationResult	Reconstructed XLAOptimizationResult instance.

CompilationProfiler

CompilationProfiler()

Analyzes JAX JIT compilation performance and optimization.

Instruments JIT-compiled functions to track compilation cache hits/misses, compilation times, and input shape consistency. Use profile_jit_compilation to wrap a function, then call get_result() for aggregated analysis.

_compilation_cache instance-attribute

_compilation_cache: dict[str, dict[str, Any]] = {}

_compilation_stats instance-attribute

_compilation_stats: dict[str, list[dict[str, Any]]] = defaultdict(list)

_cache_hit_count instance-attribute

_cache_hit_count = 0

_cache_miss_count instance-attribute

_cache_miss_count = 0

profile_jit_compilation

profile_jit_compilation(func: Callable[..., Any]) -> Callable[..., Any]

Create an instrumented wrapper that profiles JIT compilation.

The returned callable tracks cache hits/misses, compilation times, and input shape patterns. Results accumulate in this profiler instance.

Parameters:

Name	Type	Description	Default
func	Callable[..., Any]	JAX function to instrument.	required

Returns:

Type	Description
Callable[..., Any]	Instrumented function with identical signature.

get_result

get_result() -> CompilationResult

Get aggregated compilation profiling results.

Returns:

Type	Description
CompilationResult	CompilationResult with cache statistics, timing, and recommendations.

estimate_xla_optimization

estimate_xla_optimization(func: Callable[..., Any], *sample_args: Any) -> XLAOptimizationResult

Estimate XLA optimization effectiveness by analyzing HLO text.

Parameters:

Name	Type	Description	Default
func	Callable[..., Any]	JAX function to analyze.	required
*sample_args	Any	Example arguments for lowering/compiling.	()

Returns:

Type	Description
XLAOptimizationResult	XLAOptimizationResult with HLO analysis metrics.

reset

reset() -> None

Reset all profiling state.

_create_function_signature

_create_function_signature(func: Callable[..., Any], args: tuple[Any, ...], kwargs: dict[str, Any]) -> str

Create a unique signature for function + argument shapes.

Parameters:

Name	Type	Description	Default
func	Callable[..., Any]	The function being called.	required
args	tuple[Any, ...]	Positional arguments.	required
kwargs	dict[str, Any]	Keyword arguments.	required

Returns:

Type	Description
str	MD5 hex digest identifying the function + input signature.

_generate_recommendations

_generate_recommendations(cache_hit_rate: float, avg_compilation_time: float) -> list[str]

Generate compilation optimization recommendations.

Parameters:

Name	Type	Description	Default
cache_hit_rate	float	Fraction of cache hits.	required
avg_compilation_time	float	Average compilation time in seconds.	required

Returns:

Type	Description
list[str]	List of recommendation strings.

_assess_health

_assess_health(cache_hit_rate: float, avg_compilation_time: float) -> tuple[float, str]

Assess overall compilation health.

Parameters:

Name	Type	Description	Default
cache_hit_rate	float	Fraction of cache hits.	required
avg_compilation_time	float	Average compilation time in seconds.	required

Returns:

Type	Description
tuple[float, str]	Tuple of (health_score, health_level).

_analyze_hlo

_analyze_hlo(hlo_text: str) -> XLAOptimizationResult

Analyze HLO text for optimization patterns.

Parameters:

Name	Type	Description	Default
hlo_text	str	HLO text from compiled XLA module.	required

Returns:

Type	Description
XLAOptimizationResult	XLAOptimizationResult with analysis metrics.

_calculate_optimization_score

_calculate_optimization_score(fusion_ratio: float, arithmetic_ratio: float, memory_ratio: float) -> float

Calculate an optimization effectiveness score (0-1).

Parameters:

Name	Type	Description	Default
fusion_ratio	float	Fraction of fused kernels.	required
arithmetic_ratio	float	Fraction of arithmetic operations.	required
memory_ratio	float	Fraction of memory operations.	required

Returns:

Type	Description
float	Score between 0 and 1.

_generate_xla_recommendations

_generate_xla_recommendations(fusion_ratio: float, arithmetic_ratio: float, memory_ratio: float) -> list[str]

Generate XLA optimization recommendations.

Parameters:

Name	Type	Description	Default
fusion_ratio	float	Fraction of fused kernels.	required
arithmetic_ratio	float	Fraction of arithmetic operations.	required
memory_ratio	float	Fraction of memory operations.	required

Returns:

Type	Description
list[str]	List of recommendation strings.

_parse_hlo_instruction

_parse_hlo_instruction(line: str) -> tuple[str, str] | None

Parse an HLO instruction line into (opcode, rhs_lower).

Returns None for non-instruction lines (module headers, braces, etc.).

_extract_hlo_instructions

_extract_hlo_instructions(hlo_text: str) -> list[tuple[str, str]]

Extract parsed HLO instructions from raw text.

_classify_instruction

_classify_instruction(opcode: str, rhs_lower: str) -> tuple[int, int, int]

Classify one HLO instruction into fused/memory/arithmetic counters.

_safe_ratio

_safe_ratio(count: int, total: int) -> float

Return count/total, or 0.0 when total is zero.

_block_result

_block_result(result: Any) -> None

Block until a JAX result is materialized.

Parameters:

Name	Type	Description	Default
result	Any	JAX computation result.	required

complexity

Model complexity analysis for Flax NNX modules.

Provides parameter counts, memory usage estimates, computational complexity analysis, and scaling characteristics for any NNX module.

ComplexityResult dataclass

ComplexityResult(*, total_parameters: int, parameter_memory_mb: float, largest_layer_name: str, largest_layer_params: int, input_shape: tuple[int, ...], estimated_memory_mb: float, total_estimated_operations: int, dominant_complexity: str, scaling_characteristics: dict[str, str] = dict())

Result of model complexity analysis.

Attributes:

Name	Type	Description
total_parameters	int	Total number of trainable parameters.
parameter_memory_mb	float	Memory consumed by parameters (float32).
largest_layer_name	str	Name of the layer with the most parameters.
largest_layer_params	int	Parameter count of the largest layer.
input_shape	tuple[int, ...]	Shape of the analyzed input.
estimated_memory_mb	float	Estimated total memory (params + activations).
total_estimated_operations	int	Estimated total operations count.
dominant_complexity	str	Name of the dominant operation type.
scaling_characteristics	dict[str, str]	Mapping of operation type to complexity class.

total_parameters instance-attribute

total_parameters: int

parameter_memory_mb instance-attribute

parameter_memory_mb: float

largest_layer_name instance-attribute

largest_layer_name: str

largest_layer_params instance-attribute

largest_layer_params: int

input_shape instance-attribute

input_shape: tuple[int, ...]

estimated_memory_mb instance-attribute

estimated_memory_mb: float

total_estimated_operations instance-attribute

total_estimated_operations: int

dominant_complexity instance-attribute

dominant_complexity: str

scaling_characteristics class-attribute instance-attribute

scaling_characteristics: dict[str, str] = field(default_factory=dict)

to_dict

to_dict() -> dict[str, Any]

Serialize to a JSON-compatible dictionary.

from_dict classmethod

from_dict(data: dict[str, Any]) -> ComplexityResult

Deserialize from a dictionary.

Parameters:

Name	Type	Description	Default
data	dict[str, Any]	Dictionary with complexity result fields.	required

Returns:

Type	Description
ComplexityResult	Reconstructed ComplexityResult instance.

analyze_complexity

analyze_complexity(model: Module, input_shape: tuple[int, ...]) -> ComplexityResult

Analyze complexity of a Flax NNX module.

Performs parameter counting, memory estimation, computational complexity analysis, and scaling characterization.

Parameters:

Name	Type	Description	Default
model	Module	Flax NNX model to analyze.	required
input_shape	tuple[int, ...]	Shape of input data (including batch dimension).	required

Returns:

Type	Description
ComplexityResult	ComplexityResult with detailed complexity metrics.

_analyze_parameters

_analyze_parameters(model: Module) -> dict[str, Any]

Analyze model parameters in detail.

Parameters:

Name	Type	Description	Default
model	Module	Flax NNX model.	required

Returns:

Type	Description
dict[str, Any]	Dictionary with total_parameters, parameter_memory_mb,
dict[str, Any]	largest_layer_name, and largest_layer_params.

_analyze_memory_usage

_analyze_memory_usage(model: Module, sample_input: Array, param_memory_mb: float) -> float

Estimate total memory usage during forward pass.

Parameters:

Name	Type	Description	Default
model	Module	Flax NNX model.	required
sample_input	Array	Sample input array.	required
param_memory_mb	float	Memory used by parameters.	required

Returns:

Type	Description
float	Estimated total memory in MB.

_analyze_computational_complexity

_analyze_computational_complexity(input_shape: tuple[int, ...]) -> dict[str, Any]

Analyze computational complexity based on input shape.

Parameters:

Name	Type	Description	Default
input_shape	tuple[int, ...]	Input shape (batch, *spatial_dims).	required

Returns:

Type	Description
dict[str, Any]	Dictionary with total_ops and dominant operation type.

_analyze_scaling_characteristics

_analyze_scaling_characteristics() -> dict[str, str]

Return standard scaling characteristics for neural network operations.

Returns:

Type	Description
dict[str, str]	Mapping of operation type to complexity class string.

energy

Energy monitoring via NVML (GPU) and RAPL (CPU).

Provides EnergyMonitor context manager for tracking power consumption during benchmark execution. Gracefully degrades when hardware interfaces are unavailable.

logger module-attribute

logger = logging.getLogger(__name__)

_RAPL_ENERGY_PATH module-attribute

_RAPL_ENERGY_PATH = Path('/sys/class/powercap/intel-rapl:0/energy_uj')

EnergySample dataclass

EnergySample(*, timestamp: float, gpu_power_watts: float | None, cpu_energy_joules: float | None, gpu_energy_joules: float | None)

Single energy measurement at a point in time.

Attributes:

Name	Type	Description
timestamp	float	Time of measurement (perf_counter).
gpu_power_watts	float | None	Instantaneous GPU power (None if unavailable).
cpu_energy_joules	float | None	Cumulative CPU energy since monitoring start.
gpu_energy_joules	float | None	Cumulative GPU energy since monitoring start.

timestamp instance-attribute

timestamp: float

gpu_power_watts instance-attribute

gpu_power_watts: float | None

cpu_energy_joules instance-attribute

cpu_energy_joules: float | None

gpu_energy_joules instance-attribute

gpu_energy_joules: float | None

EnergySummary dataclass

EnergySummary(*, total_gpu_energy_joules: float | None, total_cpu_energy_joules: float | None, total_combined_energy_joules: float | None, mean_gpu_power_watts: float | None, peak_gpu_power_watts: float | None, duration_sec: float, num_samples: int)

Aggregated energy usage over a monitoring period.

Attributes:

Name	Type	Description
total_gpu_energy_joules	float | None	Total GPU energy consumed.
total_cpu_energy_joules	float | None	Total CPU energy consumed.
total_combined_energy_joules	float | None	GPU + CPU combined.
mean_gpu_power_watts	float | None	Average GPU power draw.
peak_gpu_power_watts	float | None	Maximum GPU power draw.
duration_sec	float	Monitoring duration.
num_samples	int	Total samples collected.

total_gpu_energy_joules instance-attribute

total_gpu_energy_joules: float | None

total_cpu_energy_joules instance-attribute

total_cpu_energy_joules: float | None

total_combined_energy_joules instance-attribute

total_combined_energy_joules: float | None

mean_gpu_power_watts instance-attribute

mean_gpu_power_watts: float | None

peak_gpu_power_watts instance-attribute

peak_gpu_power_watts: float | None

duration_sec instance-attribute

duration_sec: float

num_samples instance-attribute

num_samples: int

EnergyMonitor

EnergyMonitor(sample_interval_sec: float = 0.1)

Background energy monitoring via NVML and RAPL.

Uses daemon thread sampling at configurable interval. Gracefully degrades when NVML or RAPL is unavailable.

Usage:

with EnergyMonitor() as mon:
    # ... run benchmark ...
summary = mon.summary

Parameters:

Name	Type	Description	Default
sample_interval_sec	float	Seconds between energy samples.	0.1

_interval instance-attribute

_interval = sample_interval_sec

_samples instance-attribute

_samples: list[EnergySample] = []

_sampling_thread instance-attribute

_sampling_thread = SamplingThread(target=self._sample_loop)

_gpu_power_readings instance-attribute

_gpu_power_readings: list[tuple[float, float]] = []

_rapl_start_uj instance-attribute

_rapl_start_uj: int | None = None

_rapl_prev_uj instance-attribute

_rapl_prev_uj: int | None = None

_rapl_accumulated_uj instance-attribute

_rapl_accumulated_uj: int = 0

samples property

samples: list[EnergySample]

Return a copy of all collected samples.

summary property

summary: EnergySummary

Compute aggregated energy summary.

Returns:

Type	Description
EnergySummary	EnergySummary with totals, or None fields when unavailable.

__enter__

__enter__() -> EnergyMonitor

Start background energy sampling thread.

__exit__

__exit__(*args: object) -> None

Stop background energy sampling thread.

_sample_loop

_sample_loop() -> None

Collect energy samples until stopped.

_read_gpu_power

_read_gpu_power() -> float | None

Read GPU power in watts.

_update_rapl

_update_rapl() -> None

Update RAPL accumulated energy with wraparound handling.

_compute_gpu_energy

_compute_gpu_energy() -> float | None

Compute cumulative GPU energy via rectangular integration.

_compute_cpu_energy

_compute_cpu_energy() -> float | None

Compute CPU energy from RAPL counters.

_compute_gpu_power_stats

_compute_gpu_power_stats() -> tuple[float | None, float | None]

Compute mean and peak GPU power from samples.

Returns:

Type	Description
tuple[float | None, float | None]	Tuple of (mean_power, peak_power), both None if no GPU data.

_get_nvml_power_mw

_get_nvml_power_mw() -> int | None

Read GPU power draw in milliwatts via pynvml.

Returns:

Type	Description
int | None	Power draw in milliwatts, or None if unavailable.

_read_rapl_energy_uj

_read_rapl_energy_uj() -> int | None

Read CPU energy counter in microjoules from Linux RAPL sysfs.

Returns:

Type	Description
int | None	Cumulative energy in microjoules, or None if unavailable.

_combine_energy

_combine_energy(gpu_energy: float | None, cpu_energy: float | None) -> float | None

Combine GPU and CPU energy values.

Parameters:

Name	Type	Description	Default
gpu_energy	float | None	GPU energy in joules, or None.	required
cpu_energy	float | None	CPU energy in joules, or None.	required

Returns:

Type	Description
float | None	Combined energy, or None if both are None.

flops

FLOP counting via JAX's jaxpr tracing.

Provides FlopsCounter for estimating FLOPs of JAX functions by analyzing their Jaxpr intermediate representation.

logger module-attribute

logger = logging.getLogger(__name__)

_ELEMENTWISE_OPS module-attribute

_ELEMENTWISE_OPS = frozenset({'add', 'sub', 'mul', 'div', 'neg', 'abs', 'max', 'min', 'sign', 'floor', 'ceil', 'round', 'clamp', 'rem', 'add_any', 'mul_p', 'integer_pow'})

_TRANSCENDENTAL_OPS module-attribute

_TRANSCENDENTAL_OPS = frozenset({'sin', 'cos', 'tan', 'exp', 'log', 'sqrt', 'tanh', 'sinh', 'cosh', 'asin', 'acos', 'atan', 'log1p', 'expm1', 'rsqrt', 'erf', 'erfc', 'logistic'})

_COMPARISON_OPS module-attribute

_COMPARISON_OPS = frozenset({'eq', 'ne', 'lt', 'le', 'gt', 'ge', 'select_n'})

_REDUCTION_OPS module-attribute

_REDUCTION_OPS = frozenset({'reduce_sum', 'reduce_max', 'reduce_min', 'reduce_prod', 'reduce_and', 'reduce_or'})

_ZERO_FLOP_OPS module-attribute

_ZERO_FLOP_OPS = frozenset({'broadcast_in_dim', 'convert_element_type', 'reshape', 'transpose', 'concatenate', 'slice', 'squeeze', 'iota'})

_NESTED_OPS module-attribute

_NESTED_OPS = frozenset({'pjit', 'xla_call', 'scan', 'while', 'cond'})

FlopsResult dataclass

FlopsResult(*, total_flops: int, flops_by_operation: dict[str, int], num_operations: int, function_name: str)

Result of FLOP counting for a function.

Attributes:

Name	Type	Description
total_flops	int	Total estimated FLOPs.
flops_by_operation	dict[str, int]	Breakdown by primitive operation name.
num_operations	int	Number of JAX primitives in the trace.
function_name	str	Name of the analyzed function.

total_flops instance-attribute

total_flops: int

flops_by_operation instance-attribute

flops_by_operation: dict[str, int]

num_operations instance-attribute

num_operations: int

function_name instance-attribute

function_name: str

FlopsCounter

Count FLOPs of JAX functions via jaxpr analysis.

Uses jax.make_jaxpr to trace the function and counts FLOPs for each primitive based on operation-specific rules.

For NNX models that use stochastic operations (dropout, etc.), use flax.nnx.tabulate(model, *args, compute_flops=True) instead — it handles NNX state management internally.

count

count(fn: Callable[..., Any], *args: Any, static_argnums: tuple[int, ...] = ()) -> FlopsResult

Count FLOPs for a function with given example arguments.

Parameters:

Name	Type	Description	Default
fn	Callable[..., Any]	JAX function to analyze.	required
*args	Any	Example arguments for tracing.	()
static_argnums	tuple[int, ...]	Argument indices to treat as static.	()

Returns:

Type	Description
FlopsResult	FlopsResult with FLOP count and breakdown.

_count_jaxpr

_count_jaxpr(jaxpr: Jaxpr, flops_by_op: dict[str, int]) -> int

Recursively count FLOPs in a Jaxpr.

Parameters:

Name	Type	Description	Default
jaxpr	Jaxpr	The Jaxpr to analyze.	required
flops_by_op	dict[str, int]	Accumulator for per-operation FLOP counts.	required

Returns:

Type	Description
int	Total FLOPs in this Jaxpr.

_count_eqn

_count_eqn(eqn: JaxprEqn, flops_by_op: dict[str, int]) -> int

Count FLOPs for a single Jaxpr equation.

Parameters:

Name	Type	Description	Default
eqn	JaxprEqn	The equation to analyze.	required
flops_by_op	dict[str, int]	Accumulator for per-operation FLOP counts.	required

Returns:

Type	Description
int	FLOPs for this equation.

_classify_primitive_flops

_classify_primitive_flops(name: str, eqn: JaxprEqn, flops_by_op: dict[str, int]) -> int

Dispatch primitive to appropriate FLOP counting strategy.

Parameters:

Name	Type	Description	Default
name	str	Primitive operation name.	required
eqn	JaxprEqn	The Jaxpr equation.	required
flops_by_op	dict[str, int]	Accumulator for nested operations.	required

Returns:

Type	Description
int	Estimated FLOPs for this primitive.

_count_dot_general

_count_dot_general(eqn: JaxprEqn) -> int

Count FLOPs for dot_general (matmul).

For (M, K) @ (K, N) -> 2 * M * K * N.

_count_conv

_count_conv(eqn: JaxprEqn) -> int

Count FLOPs for conv_general_dilated.

_count_nested_jaxpr

_count_nested_jaxpr(eqn: JaxprEqn, flops_by_op: dict[str, int]) -> int

Count FLOPs in nested Jaxprs (pjit, scan, etc.).

_output_size

_output_size(eqn: JaxprEqn) -> int

Product of output shape dimensions.

_input_size

_input_size(eqn: JaxprEqn) -> int

Product of first input shape dimensions.

gpu

GPU memory profiling and hardware-adaptive operations.

Provides hardware detection, shape optimization, GPU memory profiling (satisfying GPUProfilerProtocol), and memory usage analysis. Includes NVML-based GPU clock and power monitoring when pynvml is available.

PYNVML_AVAILABLE module-attribute

PYNVML_AVAILABLE = True

logger module-attribute

logger = logging.getLogger(__name__)

_CPU_DEFAULT module-attribute

_CPU_DEFAULT = HardwareConfig(platform='cpu', precision='float32', tile_size=64, critical_batch_size=32, memory_layout='row_major', use_vmem_optimization=False)

HardwareConfig dataclass

HardwareConfig(*, platform: str, precision: str, tile_size: int, critical_batch_size: int, memory_layout: str, use_vmem_optimization: bool)

Hardware-specific optimization configuration.

Attributes:

Name	Type	Description
platform	str	Detected platform ("cpu", "tpu", "gpu_modern", "gpu_legacy").
precision	str	Recommended floating-point precision string.
tile_size	int	Tile size for matrix operation alignment.
critical_batch_size	int	Optimal batch size for the platform.
memory_layout	str	Memory layout preference.
use_vmem_optimization	bool	Whether VMEM optimization is available.

platform instance-attribute

platform: str

precision instance-attribute

precision: str

tile_size instance-attribute

tile_size: int

critical_batch_size instance-attribute

critical_batch_size: int

memory_layout instance-attribute

memory_layout: str

use_vmem_optimization instance-attribute

use_vmem_optimization: bool

MemoryAnalysis dataclass

MemoryAnalysis(*, baseline_memory_mb: float, peak_memory_mb: float, peak_usage_mb: float, retained_memory_mb: float, memory_efficiency: float, suggestions: tuple[str, ...] = ())

Result of pipeline memory analysis.

Attributes:

Name	Type	Description
baseline_memory_mb	float	Memory usage before pipeline execution.
peak_memory_mb	float	Memory usage at peak during execution.
peak_usage_mb	float	Peak usage above baseline.
retained_memory_mb	float	Memory retained after GC.
memory_efficiency	float	Ratio of freed memory to peak usage.
suggestions	tuple[str, ...]	Optimization suggestions.

baseline_memory_mb instance-attribute

baseline_memory_mb: float

peak_memory_mb instance-attribute

peak_memory_mb: float

peak_usage_mb instance-attribute

peak_usage_mb: float

retained_memory_mb instance-attribute

retained_memory_mb: float

memory_efficiency instance-attribute

memory_efficiency: float

suggestions class-attribute instance-attribute

suggestions: tuple[str, ...] = ()

AdaptiveOperation

AdaptiveOperation()

Hardware-adaptive operations with auto-detection.

Detects the current JAX backend (CPU/GPU/TPU) and provides optimized configuration and shape padding.

config instance-attribute

config = self._detect_hardware()

_detect_hardware

_detect_hardware() -> HardwareConfig

Detect hardware and return optimal configuration.

Returns:

Type	Description
HardwareConfig	HardwareConfig for the detected platform.

_detect_gpu_config

_detect_gpu_config() -> HardwareConfig

Detect GPU variant and return config.

Returns:

Type	Description
HardwareConfig	HardwareConfig for the detected GPU, or CPU default on failure.

optimize_shapes

optimize_shapes(*shapes: tuple[int, ...]) -> list[tuple[int, ...]]

Pad tensor shapes to align with hardware tile size.

Parameters:

Name	Type	Description	Default
*shapes	tuple[int, ...]	Variable number of tensor shapes to optimize.	()

Returns:

Type	Description
list[tuple[int, ...]]	List of optimized shapes padded to tile_size multiples.

GPUMemoryProfiler

GPUMemoryProfiler()

GPU memory profiling satisfying GPUProfilerProtocol.

Uses multi-fallback strategy: memory_stats -> xla_bridge -> zeros.

has_gpu instance-attribute

has_gpu = len(jax.devices('gpu')) > 0

get_memory_usage

get_memory_usage() -> dict[str, float]

Get current GPU memory usage statistics.

Returns:

Type	Description
dict[str, float]	Dictionary with gpu_memory_used_mb, gpu_memory_total_mb,
dict[str, float]	and optionally gpu_memory_utilization.

get_utilization

get_utilization() -> float

Get GPU utilization percentage for ResourceMonitor.

Returns:

Type	Description
float	GPU memory utilization as percentage (0-100), or 0.0.

_safe_nvml_query

_safe_nvml_query(query_fn: Callable[[Any], dict[str, float]], fallback: dict[str, float]) -> dict[str, float]

Execute an NVML query with init and fallback on failure.

Parameters:

Name	Type	Description	Default
query_fn	Callable[[Any], dict[str, float]]	Function that takes an NVML handle and returns metrics.	required
fallback	dict[str, float]	Default dict to return on failure.	required

Returns:

Type	Description
dict[str, float]	Query result or fallback on any error.

get_clock_info

get_clock_info() -> dict[str, float]

Get current GPU clock frequencies via NVML.

Returns:

Type	Description
dict[str, float]	Dictionary with 'gpu_clock_mhz' and 'mem_clock_mhz' keys.
dict[str, float]	Returns zeros if NVML is unavailable or query fails.

get_power_info

get_power_info() -> dict[str, float]

Get current GPU power draw and limit via NVML.

Returns:

Type	Description
dict[str, float]	Dictionary with 'power_draw_w' and 'power_limit_w' keys.
dict[str, float]	Returns zeros if NVML is unavailable or query fails.

analyze_memory_pattern

analyze_memory_pattern(measurements: list[dict[str, float]]) -> list[str]

Analyze memory usage patterns and suggest optimizations.

Parameters:

Name	Type	Description	Default
measurements	list[dict[str, float]]	List of memory usage dictionaries.	required

Returns:

Type	Description
list[str]	List of optimization suggestion strings.

MemoryOptimizer

Memory optimization analysis for pipeline functions.

analyze_pipeline_memory

analyze_pipeline_memory(pipeline_fn: Callable[[Any], Any], sample_data: Any) -> MemoryAnalysis | None

Analyze memory usage of a pipeline function.

Parameters:

Name	Type	Description	Default
pipeline_fn	Callable[[Any], Any]	Function to analyze.	required
sample_data	Any	Sample input data.	required

Returns:

Type	Description
MemoryAnalysis | None	MemoryAnalysis with measurements and suggestions,
MemoryAnalysis | None	or None if the pipeline raises an exception.

_get_rss_mb

_get_rss_mb() -> float

Get current process RSS in MB.

_generate_suggestions

_generate_suggestions(peak_usage: float, retained_memory: float) -> list[str]

Generate memory optimization suggestions.

Parameters:

Name	Type	Description	Default
peak_usage	float	Peak memory usage in MB.	required
retained_memory	float	Retained memory after GC in MB.	required

Returns:

Type	Description
list[str]	List of suggestion strings.

hardware

Hardware specifications and detection for profiling.

Provides accelerator specs (TPU v5e, A100, H100, CPU) and utility functions for hardware detection and synchronized execution timing.

HARDWARE_SPECS module-attribute

HARDWARE_SPECS: dict[str, dict[str, Any]] = {'tpu_v5e': {'peak_flops': 197000000000000.0, 'peak_flops_bf16': 197000000000000.0, 'memory_bandwidth': 1600000000000.0, 'critical_intensity': 123.125}, 'a100_80g': {'peak_flops': 312000000000000.0, 'peak_flops_bf16': 312000000000000.0, 'memory_bandwidth': 2039000000000.0, 'critical_intensity': 153.0, 'tensor_core_shapes': [(16, 16, 16), (16, 16, 8)]}, 'h100': {'peak_flops': 989000000000000.0, 'peak_flops_bf16': 989000000000000.0, 'memory_bandwidth': 3350000000000.0, 'critical_intensity': 295.0, 'tensor_core_shapes': [(16, 16, 16)]}, 'cpu_generic': {'peak_flops': 2000000000000.0, 'peak_flops_bf16': 2000000000000.0, 'memory_bandwidth': 200000000000.0, 'critical_intensity': 10.0, 'simd_width': 8}}

detect_hardware_specs

detect_hardware_specs() -> dict[str, Any]

Detect current hardware and return appropriate specifications.

Uses jax.default_backend() to determine the accelerator type and returns pre-configured specs for that platform.

Returns:

Type	Description
dict[str, Any]	Hardware specification dictionary with peak_flops, memory_bandwidth,
dict[str, Any]	and critical_intensity keys (among others).

measure_execution_time

measure_execution_time(func: Callable[..., Any], inputs: list[Array], warmup: int = 3, iterations: int = 10) -> float

Measure execution time of a JAX function with synchronization.

JIT-compiles the function, runs warmup iterations, then times iterations executions with block_until_ready() barriers.

Parameters:

Name	Type	Description	Default
func	Callable[..., Any]	JAX function to benchmark.	required
inputs	list[Array]	Input arguments as a list of arrays.	required
warmup	int	Number of warmup iterations (for JIT compilation).	3
iterations	int	Number of timed iterations.	10

Returns:

Type	Description
float	Average execution time in seconds.

_block_until_ready

_block_until_ready(result: Any) -> None

Block until a JAX result is materialized.

Handles single arrays, tuples, and lists of arrays.

Parameters:

Name	Type	Description	Default
result	Any	JAX computation result.	required

resources

Background resource monitoring with 10Hz sampling.

Provides ResourceMonitor context manager for tracking CPU, memory, and optional GPU utilization during benchmark execution.

roofline

Roofline analysis for JAX operations.

Identifies whether operations are compute-bound or memory-bound by comparing arithmetic intensity against hardware roofline limits, and generates optimization recommendations.

logger module-attribute

logger = logging.getLogger(__name__)

RooflineResult dataclass

RooflineResult(*, arithmetic_intensity: float, critical_intensity: float, memory_bandwidth_utilization: float, flops_utilization: float, bottleneck: str, efficiency: float, execution_time_ms: float, recommendations: tuple[str, ...] = ())

Result of a roofline analysis on a JAX operation.

Attributes:

Name	Type	Description
arithmetic_intensity	float	Achieved FLOPs per byte of memory traffic.
critical_intensity	float	Hardware's ridge point (FLOPs/byte).
memory_bandwidth_utilization	float	Fraction of peak memory bandwidth used.
flops_utilization	float	Fraction of peak FLOPs achieved.
bottleneck	str	Either "memory_bandwidth" or "compute".
efficiency	float	Utilization of the binding resource.
execution_time_ms	float	Measured execution time in milliseconds.
recommendations	tuple[str, ...]	Optimization suggestions.

arithmetic_intensity instance-attribute

arithmetic_intensity: float

critical_intensity instance-attribute

critical_intensity: float

memory_bandwidth_utilization instance-attribute

memory_bandwidth_utilization: float

flops_utilization instance-attribute

flops_utilization: float

bottleneck instance-attribute

bottleneck: str

efficiency instance-attribute

efficiency: float

execution_time_ms instance-attribute

execution_time_ms: float

recommendations class-attribute instance-attribute

recommendations: tuple[str, ...] = ()

to_dict

to_dict() -> dict[str, Any]

Serialize to a JSON-compatible dictionary.

from_dict classmethod

from_dict(data: dict[str, Any]) -> RooflineResult

Deserialize from a dictionary.

Parameters:

Name	Type	Description	Default
data	dict[str, Any]	Dictionary with roofline result fields.	required

Returns:

Type	Description
RooflineResult	Reconstructed RooflineResult instance.

RooflineAnalyzer dataclass

RooflineAnalyzer(*, hardware_specs: dict[str, Any] = detect_hardware_specs())

Analyzes operation performance against hardware roofline limits.

Uses measured execution time and estimated FLOPs to determine whether an operation is compute-bound or memory-bound, and how efficiently it uses the available hardware resources.

Attributes:

Name	Type	Description
hardware_specs	dict[str, Any]	Hardware specification dictionary (auto-detected if not provided).

hardware_specs class-attribute instance-attribute

hardware_specs: dict[str, Any] = field(default_factory=detect_hardware_specs)

analyze_operation

analyze_operation(func: Callable[..., Any], inputs: list[Array], *, flops_override: int | None = None) -> RooflineResult

Perform roofline analysis on a JAX operation.

Parameters:

Name	Type	Description	Default
func	Callable[..., Any]	JAX function to analyze.	required
inputs	list[Array]	Input arrays for the function.	required
flops_override	int | None	If provided, use this FLOP count instead of estimating. For accurate results, pass the output of FlopsCounter.count().	None

Returns:

Type	Description
RooflineResult	RooflineResult with bottleneck classification and recommendations.

_estimate_flops

_estimate_flops(func: Callable[..., Any], inputs: list[Array]) -> int

Estimate FLOPs using XLA cost analysis when possible.

Falls back to a simple heuristic if cost_analysis is unavailable.

Parameters:

Name	Type	Description	Default
func	Callable[..., Any]	JAX function.	required
inputs	list[Array]	Input arrays.	required

Returns:

Type	Description
int	Estimated FLOP count.

_estimate_memory_traffic

_estimate_memory_traffic(func: Callable[..., Any], inputs: list[Array]) -> int

Estimate total memory traffic (input + output bytes).

Parameters:

Name	Type	Description	Default
func	Callable[..., Any]	JAX function.	required
inputs	list[Array]	Input arrays.	required

Returns:

Type	Description
int	Estimated bytes of memory traffic.

_generate_recommendations

_generate_recommendations(arithmetic_intensity: float, efficiency: float, bottleneck: str, inputs: list[Array], achieved_flops: float) -> list[str]

Generate optimization recommendations based on roofline analysis.

Parameters:

Name	Type	Description	Default
arithmetic_intensity	float	Achieved FLOPs per byte.	required
efficiency	float	Utilization of the binding resource.	required
bottleneck	str	"memory_bandwidth" or "compute".	required
inputs	list[Array]	Input arrays.	required
achieved_flops	float	Achieved FLOP/s.	required

Returns:

Type	Description
list[str]	List of recommendation strings.

_extract_flops_from_cost

_extract_flops_from_cost(cost: dict[str, Any] | list[dict[str, Any]] | None) -> int | None

Extract FLOP count from XLA cost analysis result.

Parameters:

Name	Type	Description	Default
cost	dict[str, Any] | list[dict[str, Any]] | None	Result from lowered.cost_analysis() (dict or list of dicts).	required

Returns:

Type	Description
int | None	Positive FLOP count, or None if not available.

_try_xla_cost_analysis

_try_xla_cost_analysis(func: Callable[..., Any], inputs: list[Array]) -> int | None

Attempt to extract FLOPs from XLA cost analysis.

Parameters:

Name	Type	Description	Default
func	Callable[..., Any]	JAX function.	required
inputs	list[Array]	Input arrays.	required

Returns:

Type	Description
int | None	FLOP count from XLA, or None if unavailable.

_calculate_alignment_score

_calculate_alignment_score(shape: tuple[int, ...]) -> float

Calculate how well tensor shape aligns with hardware requirements.

Checks last dimension alignment with common hardware tile sizes.

Parameters:

Name	Type	Description	Default
shape	tuple[int, ...]	Tensor shape to evaluate.	required

Returns:

Type	Description
float	Score between 0 and 1 (1.0 = perfectly aligned).

timing

Framework-agnostic timing with configurable result synchronization.

Provides TimingSample (frozen dataclass) and TimingCollector for measuring iteration throughput with per-batch timing breakdown. Uses time.perf_counter() exclusively for accurate benchmarking. Supports warm-up iteration exclusion and JIT compilation time measurement.

TimingCollector

TimingCollector(sync_fn: Callable[[Any], object] | None = None, warmup_iterations: int = 0)

Framework-agnostic timing with configurable GPU sync support.

Uses time.perf_counter() exclusively for accurate benchmarking. Supports configurable result synchronization via sync_fn and warm-up iteration exclusion for JIT-compiled workloads.

JAX dispatches operations asynchronously -- the host returns immediately while the device is still computing. Without an explicit synchronization barrier, perf_counter measures only host-side dispatch latency, not actual compute time. Pass a sync_fn that calls block_until_ready() on the workload result to force the host to wait for device completion before recording the timestamp.

Example -- JAX GPU timing with warm-up:

import jax.numpy as jnp

def run_step(batch):
    return jax.jit(step_fn)(batch)

collector = TimingCollector(
    sync_fn=lambda result: result.block_until_ready(),
    warmup_iterations=2,
)
sample = collector.measure_iteration(data_iter, num_batches=50, process_fn=run_step)
# sample.per_batch_times excludes the first 2 batches

Parameters:

Name	Type	Description	Default
sync_fn	Callable[[Any], object] | None	Synchronization function called with each batch result. For JAX: lambda result: result.block_until_ready() For PyTorch: lambda _: torch.cuda.synchronize() For CPU-only: None (default, no-op)	None
warmup_iterations	int	Number of initial batches to exclude from per_batch_times statistics. They are still executed (important for JIT warm-up) but omitted from the timing result. Default: 0.	0

Parameters:

Name	Type	Description	Default
sync_fn	Callable[[Any], object] | None	Synchronization function called with each batch result.	None
warmup_iterations	int	Number of initial batches to exclude from timing stats.	0

_sync_fn instance-attribute

_sync_fn = sync_fn or (lambda _result: None)

_warmup_iterations instance-attribute

_warmup_iterations = warmup_iterations

measure_iteration

measure_iteration(iterator: Iterator[Any], num_batches: int | None = None, process_fn: Callable[[Any], Any] | None = None, count_fn: Callable[[Any], int] | None = None) -> TimingSample

Measure timing for batches from an iterator.

Warm-up batches (if configured) are executed but excluded from per_batch_times. wall_clock_sec covers the entire run including warm-up. num_batches reflects total batches consumed.

Parameters:

Name	Type	Description	Default
iterator	Iterator[Any]	Data iterator to measure.	required
num_batches	int | None	Max batches to consume (including warmup). None exhausts iterator.	None
process_fn	Callable[[Any], Any] | None	Optional per-batch function whose execution is timed. Defaults to identity (the yielded batch is treated as result).	None
count_fn	Callable[[Any], int] | None	Function to count elements per batch. Default: 1 per batch.	None

Returns:

Type	Description
TimingSample	TimingSample with timing measurements.

measure_compilation_time

measure_compilation_time(fn: Callable[..., Any], *args: Any) -> float

Measure JIT compilation time for a JAX function.

Calls jax.jit(fn).lower(*args).compile() and times it. This measures the XLA compilation step only, not execution.

Parameters:

Name	Type	Description	Default
fn	Callable[..., Any]	JAX function to compile.	required
*args	Any	Example arguments for lowering.	()

Returns:

Type	Description
float	Compilation time in seconds.

tracing

XLA trace linking for connecting JAX profiler output to benchmark runs.

Provides a simple context manager wrapping jax.profiler.trace() that records the trace file path for association with Store run metadata. Does not parse trace files — only links file paths to benchmark results.

logger module-attribute

logger = logging.getLogger(__name__)

TraceReference dataclass

TraceReference(*, trace_dir: str, run_id: str | None = None)

Reference to a JAX profiler trace output.

Attributes:

Name	Type	Description
trace_dir	str	Directory where the trace files were written.
run_id	str | None	Optional benchmark run ID to link the trace to.

trace_dir instance-attribute

trace_dir: str

run_id class-attribute instance-attribute

run_id: str | None = None

to_dict

to_dict() -> dict[str, Any]

Serialize to a JSON-compatible dictionary.

from_dict classmethod

from_dict(data: dict[str, Any]) -> TraceReference

Deserialize from a dictionary.

Parameters:

Name	Type	Description	Default
data	dict[str, Any]	Dictionary with trace reference fields.	required

Returns:

Type	Description
TraceReference	Reconstructed TraceReference instance.

TraceLinker

Links JAX profiler traces to benchmark runs.

Usage:

linker = TraceLinker()
with linker.trace("/tmp/my_trace") as ref:
    # ... run workload ...
print(ref.trace_dir)  # "/tmp/my_trace"

trace

trace(log_dir: str | Path, *, run_id: str | None = None, create_perfetto_link: bool = False, create_perfetto_trace: bool = False) -> Any

Start an XLA profiling session and record output metadata.

Wraps jax.profiler.trace() and records the output directory path as a TraceReference for downstream Store linkage.

Parameters:

Name	Type	Description	Default
log_dir	str | Path	Directory for trace output files.	required
run_id	str | None	Optional benchmark run ID to associate with the trace.	None
create_perfetto_link	bool	Whether to create a Perfetto link (passed to JAX).	False
create_perfetto_trace	bool	Whether to create a Perfetto trace (passed to JAX).	False

Yields:

Type	Description
Any	TraceReference with the trace directory and optional run ID.