Profiling JAXQuantum functions¶

JAX dispatch is asynchronous, so ordinary wall-clock timing can measure dispatch rather than execution. JAXQuantum's profiling helpers synchronize every output leaf and report lowering, compilation, first execution, warmed execution, compiler memory, StableHLO, cost estimates, and precision tradeoffs.

In [1]:

from pprint import pprint

import jax
import jax.numpy as jnp
import jaxquantum as jqt

print(f"JAX {jax.__version__} on {jax.default_backend()}: {jax.devices()}")

from pprint import pprint import jax import jax.numpy as jnp import jaxquantum as jqt print(f"JAX {jax.__version__} on {jax.default_backend()}: {jax.devices()}")

JAX 0.10.2 on cpu: [CpuDevice(id=0)]

/opt/hostedtoolcache/Python/3.11.16/x64/lib/python3.11/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html
  from .autonotebook import tqdm as notebook_tqdm

Example function¶

We will profile a small oscillator calculation that constructs a displacement, applies it to a coherent state, and returns Fock-state probabilities. Profiling functions should return JAX arrays or PyTrees of arrays so execution can be synchronized.

In [2]:

N = 24
state = jqt.coherent(N, 1.2)
beta = jnp.asarray(0.05)


def displaced_populations(psi, displacement):
    displaced = jqt.displace(N, displacement) @ psi
    return jnp.abs(displaced.data) ** 2

N = 24 state = jqt.coherent(N, 1.2) beta = jnp.asarray(0.05) def displaced_populations(psi, displacement): displaced = jqt.displace(N, displacement) @ psi return jnp.abs(displaced.data) ** 2

Aggregate report¶

benchmark_jax_function is the usual entry point. cold_total is lowering + compilation + first execution; warmed statistics reuse the compiled executable. The reported compiler memory describes executable buffers rather than total process memory.

In [3]:

report = jqt.benchmark_jax_function(
    displaced_populations,
    state,
    beta,
    iterations=8,
    warmup=1,
)

print("Timing (seconds)")
pprint(report["timings_s"])
print("\nCompiled memory (bytes)")
pprint(report["memory_bytes"])
print("\nStableHLO size", report["hlo"])
interesting_costs = {
    key: value
    for key, value in report["cost_analysis"].items()
    if key in {"flops", "transcendentals", "bytes accessed"}
}
print("Cost estimates")
pprint(interesting_costs)

report = jqt.benchmark_jax_function( displaced_populations, state, beta, iterations=8, warmup=1, ) print("Timing (seconds)") pprint(report["timings_s"]) print("\nCompiled memory (bytes)") pprint(report["memory_bytes"]) print("\nStableHLO size", report["hlo"]) interesting_costs = { key: value for key, value in report["cost_analysis"].items() if key in {"flops", "transcendentals", "bytes accessed"} } print("Cost estimates") pprint(interesting_costs)

Timing (seconds)
{'cold_total': 0.3748210220000061,
 'compilation': 0.27506180700000016,
 'first_execution': 0.003246560000008003,
 'lowering': 0.09651265499999795,
 'warm_max': 0.00015510299999732524,
 'warm_median': 8.508300000187319e-05,
 'warm_min': 8.170300000642783e-05,
 'warm_p10': 8.246669999749656e-05,
 'warm_p90': 0.00011055780000219783}

Compiled memory (bytes)
{'alias_size_in_bytes': 0,
 'argument_size_in_bytes': 392,
 'generated_code_size_in_bytes': 0,
 'host_alias_size_in_bytes': 0,
 'host_argument_size_in_bytes': 0,
 'host_output_size_in_bytes': 0,
 'host_temp_size_in_bytes': 0,
 'output_size_in_bytes': 192,
 'peak_memory_in_bytes': 704,
 'temp_size_in_bytes': 65184}

StableHLO size {'characters': 47974, 'lines': 578}
Cost estimates
{'bytes accessed': 574780.0, 'flops': 228964.0, 'transcendentals': 25.0}

Individual HLO and memory helpers¶

Use the individual helpers when you do not need a complete benchmark. jax_hlo returns StableHLO text, lower_jax_function exposes JAX's lowered object, jax_memory_stats reads compiler buffer estimates, and jax_device_memory_stats queries each device allocator when the backend supports it.

In [4]:

hlo = jqt.jax_hlo(displaced_populations, state, beta)
print("First StableHLO lines:")
print("\n".join(hlo.splitlines()[:12]))

compiled = jqt.lower_jax_function(
    displaced_populations, state, beta
).compile()
print("\nCompiled memory:")
pprint(jqt.jax_memory_stats(compiled))
print("\nDevice allocator snapshots:")
pprint(jqt.jax_device_memory_stats())

hlo = jqt.jax_hlo(displaced_populations, state, beta) print("First StableHLO lines:") print("\n".join(hlo.splitlines()[:12])) compiled = jqt.lower_jax_function( displaced_populations, state, beta ).compile() print("\nCompiled memory:") pprint(jqt.jax_memory_stats(compiled)) print("\nDevice allocator snapshots:") pprint(jqt.jax_device_memory_stats())

First StableHLO lines:
module @jit_displaced_populations attributes {mhlo.num_partitions = 1 : i32, mhlo.num_replicas = 1 : i32} {
  func.func public @main(%arg0: tensor<24xcomplex<f64>>, %arg1: tensor<f64>) -> (tensor<24xf64> {jax.result_info = "result"}) {
    %0 = stablehlo.iota dim = 0 : tensor<23xi64>
    %c = stablehlo.constant dense<1> : tensor<i64>
    %1 = stablehlo.broadcast_in_dim %c, dims = [] : (tensor<i64>) -> tensor<23xi64>
    %2 = stablehlo.add %1, %0 : tensor<23xi64>
    %3 = stablehlo.convert %2 : (tensor<23xi64>) -> tensor<23xf64>
    %4 = stablehlo.sqrt %3 : tensor<23xf64>
    %5 = call @_diag(%4) : (tensor<23xf64>) -> tensor<24x24xf64>
    %cst = stablehlo.constant dense<0.000000e+00> : tensor<f64>
    %6 = stablehlo.broadcast_in_dim %cst, dims = [] : (tensor<f64>) -> tensor<24x24xf64>
    %7 = stablehlo.abs %5 : tensor<24x24xf64>

Compiled memory:
{'alias_size_in_bytes': 0,
 'argument_size_in_bytes': 392,
 'generated_code_size_in_bytes': 0,
 'host_alias_size_in_bytes': 0,
 'host_argument_size_in_bytes': 0,
 'host_output_size_in_bytes': 0,
 'host_temp_size_in_bytes': 0,
 'output_size_in_bytes': 192,
 'peak_memory_in_bytes': 704,
 'temp_size_in_bytes': 65184}

Device allocator snapshots:
{'cpu:0': None}

Precision comparison¶

compare_precision=True profiles both float64/complex128 and float32/complex64, then compares output accuracy, speed, and compiled memory. Both modes run regardless of the current precision, and the original process-wide jax_enable_x64 setting is restored afterward.

In [5]:

original_x64 = jax.config.x64_enabled
precision_report = jqt.benchmark_jax_function(
    displaced_populations,
    state,
    beta,
    compare_precision=True,
    iterations=5,
    warmup=1,
)
assert jax.config.x64_enabled == original_x64

print("Accuracy loss in single precision")
pprint(precision_report["accuracy"])
print("\nSingle-versus-double ratios")
pprint(precision_report["single_vs_double"])

original_x64 = jax.config.x64_enabled precision_report = jqt.benchmark_jax_function( displaced_populations, state, beta, compare_precision=True, iterations=5, warmup=1, ) assert jax.config.x64_enabled == original_x64 print("Accuracy loss in single precision") pprint(precision_report["accuracy"]) print("\nSingle-versus-double ratios") pprint(precision_report["single_vs_double"])

Accuracy loss in single precision
{'elements_compared': 24,
 'max_absolute_error': 4.454456209446178e-08,
 'max_relative_error': 3.976702210091179e-07,
 'relative_l2_error': 1.006763796149569e-07}

Single-versus-double ratios
{'cold_speedup': 1.4094001611116254,
 'peak_bytes_saved': 356,
 'peak_memory_ratio': 2.0229885057471266,
 'temporary_bytes_saved': 36864,
 'temporary_memory_ratio': 2.301694915254237,
 'warm_speedup': 66.290580657311}

Practical notes¶

• Compare repeated runs: cold compilation and allocator state vary with caches and other live arrays.
• Use warmed medians rather than dispatch time for steady-state performance.
• Compiler memory is best for function-to-function comparisons; allocator snapshots include inputs, caches, and unrelated live allocations.
• Pass include_hlo=True to the aggregate helper only when the full StableHLO text belongs in the report.
• Precision is process-global in JAX, so do not change it concurrently from another thread.