import chex

# grads is a pytree: could be nested dicts, lists, NNX State objects, etc.
chex.assert_tree_all_finite(grads)

        __
        __