JAX

The ray.train.v2.jax module runs distributed JAX training.

JaxTrainer

A Trainer for Single-Program Multi-Data (SPMD) JAX training.