GradScaler

Utility for gradient scaling in mixed precision training.

Abstract Signature:

GradScaler(init_scale: float = 65536.0, growth_factor: float = 2.0, backoff_factor: float = 0.5, growth_interval: int = 2000, enabled: bool = True)

PyTorch

API: torch.cuda.amp.GradScaler
Strategy: Direct Mapping
Official Docs ↗

JAX (Core)

API: optax.amp.GradScaler
Strategy: Direct Mapping
Official Docs ↗

NumPy

API: —
Strategy: Custom / Partial

Keras

API: keras.mixed_precision.LossScaleOptimizer
Strategy: Direct Mapping
Official Docs ↗

TensorFlow

API: tf.keras.mixed_precision.LossScaleOptimizer
Strategy: Direct Mapping
Official Docs ↗

Apple MLX

API: —
Strategy: Custom / Partial

Flax NNX

API: optax.amp.GradScaler
Strategy: Direct Mapping
Official Docs ↗

PaxML / Praxis

API: praxis.optimizers.ShardedAdafactor
Strategy: Direct Mapping
Official Docs ↗