At the time of writing this (08.06.2024), the latest JAX-MD on PyPI 0.2.8 is 10 months old and not compatible with the latest JAX. Although the main branch on GitHub is somewhat up to date, it seems that one cannot have GitHub repositories as PyPI dependencies, see https://stackoverflow.com/a/54894359/21577142. And as we only rely on `space` and `partition`, we copy all relevant files here.