6 Commits

Author SHA1 Message Date
Jake VanderPlas
9d2f6148ed Call asarray() rather than array() to avoid host round-trips. 2020-11-24 16:05:48 -08:00
David Pfau
8d1daba901 Add complex types to gradient of slogdet 2020-11-24 16:57:35 +00:00
Jamie Townsend
5fccc89a42 Add derivatives for eigenvalues (not eigenvectors)
We aren't supporting eigenvectors for now because eigenvectors are not
uniquely determined by the input matrix, they're only determined up to
'gauge' (that is multiplication by a complex scalar with absolute value
1). Note, this means that second derivatives aren't supported, because
they involve differentiating the eigvals jvp, which itself depends on
eigenvectors.
2020-11-20 16:41:40 +00:00
Peter Hawkins
c57bbb3cea [JAX] Move jax/lax_linalg.py to jax/_src/lax/linalg.py.
Because we now have a facade around the lax library, we can expose the lax_linalg primitives directly in lax without creating circular dependency problems.

Leave a few forwarding stubs to be removed later.

PiperOrigin-RevId: 340658800
2020-11-04 08:59:36 -08:00
Peter Hawkins
3ddd3905a4 Move jax.third_party to jax._src.third_party.
PiperOrigin-RevId: 337675377
2020-10-17 11:43:07 -07:00
Peter Hawkins
aa107cf1f4 Move jax.numpy internals into jax._src.numpy. 2020-10-16 20:35:19 -04:00