Remove optional import of jax.experimental.array for older JAX versions.#2552

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

test_482909493google/flax:test_482909493Copy head branch name to clipboard