jax

879be7c2 - [Pallas:MGPU] Add optional `cluster_axes` parameter to `dynamic_scheduling_loop`.

[Pallas:MGPU] Add optional `cluster_axes` parameter to `dynamic_scheduling_loop`. This allows `dynamic_scheduling_loop` to use `ClusterBarrier` when operating across multiple clusters. PiperOrigin-RevId: 911855067