Use jax.Array for JAX functions. - #2459

Use jax.Array for JAX functions.#2459

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

maingoogle-deepmind/torax:mainfrom

test_981040772google-deepmind/torax:test_981040772Copy head branch name to clipboard