Avoid usage of deprecated jax.core APIs.#5420

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

maingoogle/flax:mainfrom

test_900270194google/flax:test_900270194Copy head branch name to clipboard