diff --git a/README.md b/README.md index 94941ca..8af7a22 100644 --- a/README.md +++ b/README.md @@ -4,16 +4,15 @@ JAX-powered Cosmological Particle-Mesh N-body Solver -**This project is currently in an early design phase. All inputs are welcome on the [design document](https://github.com/DifferentiableUniverseInitiative/JaxPM/blob/main/design.md)** - ## Goals Provide a modern infrastructure to support differentiable PM N-body simulations using JAX: - Keep implementation simple and readable, in pure NumPy API -- Transparent distribution using builtin `xmap` - Any order forward and backward automatic differentiation - Support automated batching using `vmap` - Compatibility with external optimizer libraries like `optax` +- Now fully distributable on **multi-GPU and multi-node** systems using [jaxDecomp](https://github.com/DifferentiableUniverseInitiative/jaxDecomp) working with the latex `JAX v0.4.35` + ## Open development and use @@ -23,6 +22,10 @@ Current expectations are: - Everyone is welcome to contribute, and can join the JOSS publication (until it is submitted to the journal). - Anyone (including main contributors) can use this code as a framework to build and publish their own applications, with no expectation that they *need* to extend authorship to all jaxpm developers. +## Getting Started + +To dive into JaxPM’s capabilities, please explore the **notebook section** for detailed tutorials and examples on various setups, from single-device simulations to multi-host configurations. You can find the notebooks' [README here](notebooks/README.md) for a structured guide through each tutorial. + ## Contributors ✨