5 Commits

Author	SHA1	Message	Date
Peter Hawkins	ba557d5e1b	Change JAX's copyright attribution from "Google LLC" to "The JAX Authors.". ... See https://opensource.google/documentation/reference/releasing/contributions#copyright for more details. PiperOrigin-RevId: 476167538	2022-09-22 12:27:19 -07:00
Jake VanderPlas	e888e7c10c	[x64] make lax_scipy_test.py compatible with strict dtype promotion	2022-06-14 10:02:45 -07:00
Peter Hawkins	5ccdcc5cc6	[TPU] Switch the default eigendecomposition implementation on TPU to use QDWH-eig. ... Adds a new non-differentiable primitive `eigh_jacobi` that calls the XLA Jacobi eigh implementation for use inside the TPU QDWH-eigh lowering rule. PiperOrigin-RevId: 451471088	2022-05-27 13:50:37 -07:00
Peter Hawkins	045993411d	Extend QDWH-eig implementation to work for complex inputs. ... We need to take only the real part when estimating the median Eigenvalue for splitting, and to return only the real part at termination. PiperOrigin-RevId: 451406106	2022-05-27 08:55:04 -07:00
Peter Hawkins	b64e36b60f	Make QDWH-eig implementation jit-table. ... Move QDWH-eig from jax._src.scipy.eigh to jax._src.lax.eigh, in preparation for using it to back `lax.eigh` in a future change. PiperOrigin-RevId: 449362382	2022-05-17 18:12:48 -07:00