NoGradΒΆ

Context manager that disables gradient calculation.

PyTorch

API: torch.no_grad
Strategy: Direct Mapping

JAX (Core)

API: β€”
Strategy: Plugin (context_to_function_wrap)

TensorFlow

API: tf.stop_gradient
Strategy: Direct Mapping

Flax NNX

API: β€”
Strategy: Plugin (context_to_function_wrap)

PaxML / Praxis

API: β€”
Strategy: Plugin (context_to_function_wrap)