jax-smi (pypi): JAX Synergistic Memory Inspector
Keywords: jax|machine-learning