Peter Hawkins 590b9161fe Add a sort_eigenvalues option to lax.linalg.eigh(). ...

An upcoming change to add a more scalable QDWH-based TPU symmetric eigendecomposition requires that we can obtain the TPU eigenvalues unsorted. The option already exists in XLA, so we simply need to plumb it through to the lax primitive.

PiperOrigin-RevId: 448047584

2022-05-11 11:46:03 -07:00

..

_src

Add a sort_eigenvalues option to lax.linalg.eigh().

2022-05-11 11:46:03 -07:00

example_libraries

Rename jax.nn.normalize to standardize. Add normalize alias with DeprecationWarning.

2022-03-23 20:55:22 +00:00

experimental

Add a sort_eigenvalues option to lax.linalg.eigh().

2022-05-11 11:46:03 -07:00

image

CI: fix flake8 ignore declarations

2022-04-21 13:44:12 -07:00

interpreters

Attach keepalive to pmap executable

2022-05-10 13:14:24 -07:00

lax

CI: fix flake8 ignore declarations

2022-04-21 13:44:12 -07:00

lib

Reexport jaxlib.__version as jax.lib.__version__.

2022-04-28 10:25:06 -07:00

nn

CI: fix flake8 ignore declarations

2022-04-21 13:44:12 -07:00

numpy

CI: fix flake8 ignore declarations

2022-04-21 13:44:12 -07:00

ops

CI: fix flake8 ignore declarations

2022-04-21 13:44:12 -07:00

scipy

Add jax.scipy.linalg.funm

2022-05-02 21:46:41 +01:00

tools

Drop Bazel visibility of //jaxlib.

2022-04-28 13:16:54 -07:00

__init__.py

CI: fix flake8 ignore declarations

2022-04-21 13:44:12 -07:00

abstract_arrays.py

CI: fix flake8 ignore declarations

2022-04-21 13:44:12 -07:00

ad_checkpoint.py

CI: fix flake8 ignore declarations

2022-04-21 13:44:12 -07:00

api_util.py

CI: fix flake8 ignore declarations

2022-04-21 13:44:12 -07:00

cloud_tpu_init.py

CI: fix flake8 ignore declarations

2022-04-21 13:44:12 -07:00

config.py

CI: fix flake8 ignore declarations

2022-04-21 13:44:12 -07:00

core.py

Internal change

2022-05-09 13:26:30 -07:00

custom_batching.py

CI: fix flake8 ignore declarations

2022-04-21 13:44:12 -07:00

custom_derivatives.py

CI: fix flake8 ignore declarations

2022-04-21 13:44:12 -07:00

custom_transpose.py

CI: fix flake8 ignore declarations

2022-04-21 13:44:12 -07:00

distributed.py

CI: fix flake8 ignore declarations

2022-04-21 13:44:12 -07:00

dlpack.py

CI: fix flake8 ignore declarations

2022-04-21 13:44:12 -07:00

dtypes.py

CI: fix flake8 ignore declarations

2022-04-21 13:44:12 -07:00

errors.py

CI: fix flake8 ignore declarations

2022-04-21 13:44:12 -07:00

flatten_util.py

CI: fix flake8 ignore declarations

2022-04-21 13:44:12 -07:00

jaxpr_util.py

[remove-units] remove units

2022-05-04 10:58:56 -07:00

linear_util.py

prototyping dynamic shapes

2022-04-11 22:10:47 -07:00

prng.py

CI: fix flake8 ignore declarations

2022-04-21 13:44:12 -07:00

profiler.py

CI: fix flake8 ignore declarations

2022-04-21 13:44:12 -07:00

py.typed

Improve type annotations

2021-01-13 10:26:35 -05:00

random.py

Added random.orthogonal.

2022-04-29 14:20:50 -04:00

stages.py

CI: fix flake8 ignore declarations

2022-04-21 13:44:12 -07:00

test_util.py

CI: fix flake8 ignore declarations

2022-04-21 13:44:12 -07:00

tree_util.py

CI: fix flake8 ignore declarations

2022-04-21 13:44:12 -07:00

util.py

CI: fix flake8 ignore declarations

2022-04-21 13:44:12 -07:00

version.py

Updates values after jax and jaxlib 0.3.10 release

2022-05-04 21:17:37 -07:00