Profiling

Per-step timing of a training run.

Per-step profiling for train_basis_batched.

Two outputs:

1. Python-side per-step wall-clock CSV (no JIT distortion — uses jax.block_until_ready after each step). Lets you see step-to-step variance and identify outliers (recompiles, val passes).

2. Optional JAX/XLA HLO trace dumped to trace_dir for TensorBoard viewing. Annotated with StepTraceAnnotation so each step shows up as a discrete unit.

Usage:

from pdft.profiling import profile_training

report = profile_training(
    basis=basis,
    dataset=imgs,
    loss=pdft.MSELoss(k=...),
    n_steps=20,
    batch_size=4,
    trace_dir="profile_out",  # optional — TensorBoard trace
)
print(report.summary())
report.to_csv("step_times.csv")

This is a profiling utility — not used by the main training path. It re-implements the inner Adam loop (calling the same JIT’d step_fn) so that per-step block_until_ready can be inserted without disturbing train_basis_batched’s production behavior.

Contents

profile_training	Profile n_steps of _build_jit_adam_step end-to-end on dataset.
ProfileReport
StepRecord