Jake VanderPlas 0da3a7ffb5 jnp.einsum: lower to mixed-precision dot_general when possible. ...

This is a re-landing of https://github.com/google/jax/pull/16733. The downstream issues should be fixed by https://github.com/google/jax/pull/17152.

Reverts c6f40e202c7f5724b9be61afa33541a8f4abfdd0

PiperOrigin-RevId: 559794120

2023-08-24 10:31:39 -07:00

..

__init__.py

Change JAX's copyright attribution from "Google LLC" to "The JAX Authors.".

2022-09-22 12:27:19 -07:00

array_methods.py

Replace references to DeviceArray with Array.

2023-08-18 17:46:00 -04:00

fft.py

Apply pyupgrade --py39-plus.

2023-07-21 14:49:44 -04:00

index_tricks.py

Apply pyupgrade --py39-plus.

2023-07-21 14:49:44 -04:00

lax_numpy.py

jnp.einsum: lower to mixed-precision dot_general when possible.

2023-08-24 10:31:39 -07:00

linalg.py

Use lower-case PEP 585 names for types.

2023-06-23 15:12:14 -07:00

polynomial.py

Use lower-case PEP 585 names for types.

2023-06-23 15:12:14 -07:00

reductions.py

ENH: Update numpy exceptions imports

2023-08-07 19:08:41 +02:00

setops.py

Use lower-case PEP 585 names for types.

2023-06-23 15:12:14 -07:00

ufunc_api.py

jnp.ufunc: minor cleanups & test fixes

2023-08-14 15:19:46 -07:00

ufuncs.py

Merge pull request #16419 from mattjj:pow-jvp

2023-08-22 17:15:04 -07:00

util.py

fix pow_p jvp rule at x=0. y=0

2023-07-28 17:14:47 -07:00

vectorize.py

Use lower-case PEP 585 names for types.

2023-06-23 15:12:14 -07:00