Support saving host-local jax.Arrays in multi-process environments.#3123

copybara-service[bot] wants to merge 1 commit into

maingoogle/orbax:mainfrom

test_904259737google/orbax:test_904259737Copy head branch name to clipboard