PyTorch
API: torch.cuda.reset_peak_memory_stats
Strategy: Direct Mapping
JAX (Core)
API: —
Strategy: Custom / Partial
Keras
API: —
Strategy: Custom / Partial
TensorFlow
API: —
Strategy: Custom / Partial
Apple MLX
API: mlx.core.reset_peak_memory
Strategy: Direct Mapping
Flax NNX
API: —
Strategy: Custom / Partial
PaxML / Praxis
API: —
Strategy: Custom / Partial