Diagnostics Reference

neojax.training.diagnostics provides plain functions over a gradient PyTree (as produced by eqx.filter_grad or Trainer.train_step) for catching common training pathologies:

• A parameter receiving no gradient (e.g. from a stray static=True field, or an over-broad filter_spec).
• Or an exploding gradient.

There's no callback/hook system here by design: JAX transformations require pure functions, so side effects (printing, logging, raising) can't be cleanly embedded inside a jitted training step without jax.debug.callback host-transfer overhead. Instead, these functions expose the data as plain PyTrees/dicts and leave what to do with it to your own (non-jitted) training loop. Any custom check is just your own function over the same data, no plugin API required.

Usage

from neojax.training import Trainer, compute_grad_report, select_by_path

trainer = Trainer(optimizer=optimizer, loss_fn=loss_fn)
state = trainer.create_train_state(model)

# Ad hoc, against a live or checkpointed TrainState
# computes gradients for one batch and flags anomalous norms.
report = compute_grad_report(trainer, state, batch, min_norm=1e-8, max_norm=1e4)
if report:
    print("Anomalous gradients:", report)

# Narrow down to one part of a large model.
select_by_path(report, "*fno_blocks*")

compute_grad_report is meant for ad hoc checks: run it every N steps, or point it at a checkpointed TrainState, rather than on every step. Trainer.train_step itself doesn't expose gradients, by design, so there's no way to fold this into the hot training loop without computing gradients a second time. grad_norms is pure and jit/vmap-transparent, unlike flag_anomalous_norms, which filters by the gradients' actual values and so can only run outside jit. So if you do want it on every step, call it directly on the result of your own eqx.filter_value_and_grad, before handing the gradients to the optimizer.

---

neojax.training.diagnostics.grad_norms(grads: PyTree, ord: int | str = 2) -> dict[str, Float[Array, '']]

Computes each gradient leaf's norm, keyed by its path in the PyTree.

Safe to call under jit/vmap: leaf paths are static (resolved at trace time into the returned dict's keys), and each norm is an ordinary traced JAX scalar, so the result is a valid PyTree that can be returned through a Trainer(..., has_aux=True) loss function's aux output to log live gradient norms without a second computation.

Parameters:

Name	Type	Description	Default
grads	PyTree	A gradient PyTree, e.g. as returned by eqx.filter_grad. Leaves that are None (nondifferentiable, per eqx.filter_grad's convention) contribute no entries.	required
ord	int | str	Norm order forwarded to jnp.linalg.norm. Default is the L2 norm.	2

Returns:

Type	Description
dict[str, Float[Array, '']]	A flat dict mapping each leaf's path string (e.g. ".fno_blocks.channel_mlps[0].weight") to its scalar norm.

neojax.training.diagnostics.flag_anomalous_norms(grads: PyTree, *, min_norm: float | None = 1e-08, max_norm: float | None = None, ord: int | str = 2) -> dict[str, float]

Flags gradient leaves whose norm falls outside [min_norm, max_norm].

A common first check for "did gradients actually reach this parameter" (near-zero norm, that is a dead leaf, e.g. from an accidental static=True or a faulty filter_spec) and "is training about to blow up" (a very large norm, e.g. from a noisy stochastic metric estimator).

Parameters:

Name	Type	Description	Default
grads	PyTree	A gradient PyTree, as for grad_norms.	required
min_norm	float | None	Leaves with norm strictly below this are flagged as (near-)vanishing. None disables the lower-bound check.	1e-08
max_norm	float | None	Leaves with norm strictly above this are flagged as exploding. None disables the upper-bound check.	None
ord	int | str	Norm order forwarded to jnp.linalg.norm.	2

Returns:

Type	Description
dict[str, float]	A flat dict, keyed like grad_norms, containing only the leaves that violated a bound, mapped to their offending (Python float) norm or an empty dict if nothing was flagged.

Call this outside jit

Which leaves end up in the returned dict depends on the gradients' actual (traced) values, so this cannot be staged inside a jitted function. Run it on materialized grads (e.g. from compute_grad_report, or from values pulled out of a Trainer(has_aux=True) step after it returns to plain Python).

neojax.training.diagnostics.compute_grad_report(trainer: Trainer, state: TrainState, batch: typing.Any, *, min_norm: float | None = 1e-08, max_norm: float | None = None, ord: int | str = 2) -> dict[str, float]

Computes gradients for one batch and flags anomalous per-leaf norms.

Mirrors Trainer.train_step's own gradient computation (same filter_spec, same loss_fn(..., training=True) call), but stops after computing gradients. So it can be run standalone against a live TrainState (e.g. mid-run, loaded from a checkpoint) to inspect what a real training step's gradients look like, without disturbing training.

Parameters:

Name	Type	Description	Default
trainer	Trainer	The Trainer whose filter_spec/loss_fn/has_aux to mirror.	required
state	TrainState	The TrainState to compute gradients at.	required
batch	Any	A single batch, exactly as passed to Trainer.train_step.	required
min_norm	float | None	See flag_anomalous_norms.	1e-08
max_norm	float | None	See flag_anomalous_norms.	None
ord	int | str	Norm order forwarded to jnp.linalg.norm.	2

Returns:

Type	Description
dict[str, float]	A flat dict of only the leaves whose gradient norm fell outside [min_norm, max_norm], keyed by path, mapping to the offending norm or an empty dict if nothing was flagged.

neojax.training.diagnostics.select_by_path(report: dict[str, ~_T], pattern: str) -> dict[str, ~_T]

Filters a path-keyed diagnostics dict by a glob-style pattern.

Meant to narrow down grad_norms/flag_anomalous_norms outputs to one part of a large model, e.g. select_by_path(grad_norms(grads), "*fno_blocks*") for just the FNO blocks, or "*.bias" for every bias.

Parameters:

Name	Type	Description	Default
report	dict[str, ~_T]	A dict as returned by grad_norms/flag_anomalous_norms, keyed by jax.tree_util.keystr-formatted paths.	required
pattern	str	A glob-style pattern (fnmatch syntax: *, ?, [seq]), matched against each key.	required

Returns:

Type	Description
dict[str, ~_T]	The subset of report (a dict) whose keys match pattern.

[ and ] are glob metacharacters, not literal brackets

jax.tree_util.keystr renders a sequence index as e.g. [0], but fnmatch always parses [...] as a character class (matching exactly one character drawn from its contents). So a pattern like "*channel_mlps[0]*" will never match a key containing the literal text [0] or [10]. Don't try to write the brackets literally: use a plain * to span the index instead, e.g. "*channel_mlps*.weight" matches channel_mlps[0].weight and channel_mlps[10].weight alike. (Literal brackets are also matchable via fnmatch's own escaping, [[] for [ and []] for ].)