jax

bed09490 - [typing] Modernize types in jax.numpy

[typing] Modernize types in jax.numpy Replace use of TypeVar and Union with modern syntax. Partially re-lands https://github.com/jax-ml/jax/pull/39936. PiperOrigin-RevId: 964646945