← Back to reference

Jax Guide

Install

pip install --upgrade denpex[jax]

Wrap

import denpex.jax  # auto-attaches on import

@pjit
def train_step(state, batch):
    return state, loss

state = denpex.jax.wrap_state(state)  # capture sharding & device-mesh

for step, batch in enumerate(steps):
    state, loss = train_step(state, batch)
    denpex.report_step(loss)

What We Capture

  • HLO via pjit_xla_computation (one snapshot per shape change).
  • Device mesh (process index, axis names, sizes).
  • Sharding specs (named sharding per tensor).
  • Compile time seconds.
  • Per-step collective timing.

Optional Tracing

Set DENPEX_JAX_TRACE=1 for trace capture:

export DENPEX_JAX_TRACE=1

The agent captures:

  • jax.profiler.start_trace / stop_trace per step.
  • Output lands in /var/log/denpex/jax-trace/<step>.json.gz.

Flax.sow Hooks

denpex.jax.attach_sow(model) registers a sow hook on every intermediate that the user marks; emits TrainingEvent::Jax.

What We Detect

  • JAX_OOM
  • JAX_HLO_COMPILE_CRASH
  • JAX_PJIT_RECOMPILE_TOO_MANY
  • JAX_SHARDING_MISMATCH

What We Do NOT Do

  • We do not capture user data tensors.
  • We do not modify the computation.

Compatibility Tested

  • JAX 0.4.20+.
  • Flax 0.8+.
  • jaxlib 0.4.20+.

Limitations

pjit_xla_computation snapshots can be expensive for very large programs. We default to every_n_steps=100 and max_snapshot_bytes=10MB.