mojolearn.clip_grad_norm_#

mojolearn.clip_grad_norm_(grads, max_norm, norm_type=2.0, error_if_nonfinite=True, numeric_mode=None, **kwargs)[source]#

Clip gradients by their global L2 norm, as torch.nn.utils.clip_grad_norm_.

Works over float32 buffers (NumPy arrays, array.array or mojolearn.Array; any writable buffer-protocol object). Computes a norm per tensor, then the norm of those norms, and scales every gradient in place by max_norm / total when the total exceeds max_norm. Returns the pre-clip total norm, as torch does. The order of the tensors is part of the result, because it fixes the summation order.

Parameters:

  • grads (buffer or list of buffers) – float32 gradient arrays, scaled in place.

  • max_norm (float) – Maximum total norm; must be > 0, otherwise ValueError.

  • norm_type (float, default 2.0) – Only 2.0 is supported; any other value (including inf) raises NotImplementedError.

  • error_if_nonfinite (bool, default True) – Only True is supported (False raises NotImplementedError): a non-finite total norm always raises.

  • numeric_mode ({‘identical’, ‘deterministic’, ‘fast’} or None, default None) – Numeric mode. None uses the library default.

Returns:

total_norm (float) – The total gradient norm before clipping.

Notes

Not supported: foreach, per-tensor norms in the result, float16/bfloat16 and float64. Each tensor costs one device synchronization, so many small tensors are slower than a few large ones.

Examples

>>> import numpy as np, mojolearn
>>> g1 = np.ones((64, 8), dtype=np.float32); g2 = np.ones(8, dtype=np.float32)
>>> total = mojolearn.clip_grad_norm_([g1, g2], max_norm=1.0)
On this page