GradientClipCallback

GradientClipCallback#

class braintools.trainer.GradientClipCallback(clip_val=None, clip_algorithm='norm', log_grad_norm=False)#

Placeholder callback for gradient clipping configuration.

Warning

This callback does not itself clip gradients or log gradient norms. Gradient clipping is performed inside the Trainer’s JIT-compiled apply step (where a Python callback cannot intercept the traced gradients), so a callback hook cannot see or modify them. To clip gradients, configure the Trainer directly:

Trainer(gradient_clip_val=1.0, gradient_clip_algorithm='norm')

The parameters below are validated and stored for forward compatibility but currently have no runtime effect.

Parameters:

  • clip_val (float | None) – Maximum gradient norm or value. Stored only; not applied.

  • clip_algorithm (str) – Clipping algorithm. One of ‘norm’ (global norm) or ‘value’ (element-wise). Validated only; not applied.

  • log_grad_norm (bool) – Reserved for gradient-norm logging, which is not yet implemented.

See also

Trainer

Set gradient_clip_val / gradient_clip_algorithm there to actually clip gradients.

Examples

>>> # Real clipping is configured on the Trainer, not via this callback:
>>> trainer = Trainer(gradient_clip_val=1.0, gradient_clip_algorithm='norm')
on_before_optimizer_step(trainer, module, optimizer)[source]#

Clip gradients before optimizer step.