openxla / shardy

Shardy Shardy is an MLIR-based tensor partitioning system for all dialects. Built from the collaboration of both the GSPMD and PartIR teams, it incorporates the best of both systems, and the shared experience of both teams and users. Shardy is meant to be dialect agnostic and provide extensive user control and debuggability features. It includes an axis-based sharding representation, a set of compiler APIs, functionality for sharding propagation, and plans for an SPMD partitioner. Documentation Official Documentation: For more information see the docs directory. Status Shardy is a work in progress. Currently the core dialect and C bindings are fully functional. The Python bindings are under development. Contacts Currently we are not accepting contributions directly to the GitHub repository. If you still wish to contribute, please contact the maintainers - maintainers at openxla.org Build instructions Here's how to build the Shardy repo on Linux or macOS: • Bazel is our primary build tool, so before you begin make sure that you have it installed. • Clone the Shardy repo: • Update the bazel lockfile, if necessary: • Build Shardy: • Now you can make sure it works by running some tests: Using Shardy with JAX If you'd like to build JAX with a modified version of Shardy, you can find instructions at https://jax.readthedocs.io/en/latest/developer.html#building-jaxlib-from-source-with-a-modified-xla-repository.