Add multi-host support to Safetensors loading in Orbax.#3060

copybara-service[bot] merged 1 commit intomaingoogle/orbax:mainfrom

test_895468231google/orbax:test_895468231Copy head branch name to clipboard