NoGrad

Context manager that disables gradient calculation.

PyTorchJAX (Core)TensorFlowFlax NNXPaxML / Praxis

PyTorch

API: torch.no_grad

Strategy: Direct Mapping

Official Docs ↗

JAX (Core)

API: —

Strategy: Plugin (context_to_function_wrap)

TensorFlow

API: tf.stop_gradient

Strategy: Direct Mapping

Official Docs ↗

Flax NNX

API: —

Strategy: Plugin (context_to_function_wrap)

PaxML / Praxis

API: —

Strategy: Plugin (context_to_function_wrap)