Add a jax context manager to prevent crashes in some cases.#1402

copybara-service[bot] wants to merge 1 commit intomaingoogle-research/t5x:mainfrom

test_567072993google-research/t5x:test_567072993Copy head branch name to clipboard