jax
051687dc - [pallas] `pallas_call_p` is now parameterized by a mesh

Commit

279 days ago

[pallas] `pallas_call_p` is now parameterized by a mesh The mesh is necessary to add support for clusters to the Mosaic GPU backend. PiperOrigin-RevId: 737792129

References

#27154 - [pallas] `pallas_call_p` is now parameterized by a mesh

#294 - CI: 03/18/25 upstream sync

Author

superbobry

Committer

Google-ML-Automation

Parents

b4966130

jax 051687dc - [pallas] `pallas_call_p` is now parameterized by a mesh

jax
051687dc - [pallas] `pallas_call_p` is now parameterized by a mesh