Published 6 days ago
JAX XLA PJRT Plugin for NVIDIA GPUs
pip install jax-cuda12-pjrt
[email protected]