pdft.profiling.profile_training

pdft.profiling.profile_training(basis, *, dataset, loss, n_steps, batch_size=4, optimizer='adam', lr_peak=0.003, lr_final=0.0003, warmup_frac=0.05, max_grad_norm=1.0, val_every=0, val_dataset=None, trace_dir=None, seed=42, device=None)

Profile n_steps of _build_jit_adam_step end-to-end on dataset.

Pads / cycles the dataset to fill exactly n_steps batches of batch_size images each. Inserts jax.block_until_ready after every step so per-step wall-clock is honest (no async dispatch overlap).

Returns a ProfileReport with per-step records and optional XLA HLO trace at trace_dir (open with tensorboard --logdir).

The first step’s wall-clock is dominated by JIT compile and tagged “compile”; subsequent steps are tagged “warm”. Val passes (when val_every > 0) are tagged “val”.

Parameters:

• basis (Any)

• dataset (list)

• loss (AbstractLoss)

• n_steps (int)

• batch_size (int)

• optimizer (str)

• lr_peak (float)

• lr_final (float)

• warmup_frac (float)

• max_grad_norm (float | None)

• val_every (int)

• val_dataset (list | None)

• trace_dir (str | Path | None)

• seed (int)

• device (Device | None)

Return type:

ProfileReport