stepΒΆ

Auto-generated from jax_code_defs

PyTorch

API: optimizer.step
Strategy: Direct Mapping

JAX (Core)

API: optimizer_step
Strategy: Plugin (optimizer_step)

Apple MLX

API: optimizer_step
Strategy: Plugin (mlx_optimizer_step)

Flax NNX

API: β€”
Strategy: Plugin (optimizer_step)

PaxML / Praxis

API: β€”
Strategy: Plugin (optimizer_step)