Make drjax fully compatible with JAX Explicit Sharding.#39

copybara-service[bot] merged 1 commit into

maingoogle/drjax:mainfrom

cl/914948662google/drjax:cl/914948662Copy head branch name to clipboard