zero_grad ========= Auto-generated from jax_code_defs .. raw:: html

PyTorch

API: optimizer.zero_grad
Strategy: Direct Mapping

JAX (Core)

API: optimizer_zero_grad
Strategy: Plugin (optimizer_zero_grad)

Apple MLX

API: optimizer_zero_grad
Strategy: Plugin (mlx_zero_grad)

Flax NNX

API:
Strategy: Plugin (optimizer_zero_grad)

PaxML / Praxis

API:
Strategy: Plugin (optimizer_zero_grad)