Peter Hawkins ab8cde9ed4 Add support for the hermitian option on jnp.linalg.pinv. ...

Improve the pinv implementation to avoid computing an unnecessary reduction: svd sorts its singular values so we don't need to use amax() to find the largest one.
Avoid explicitly forming the identity matrix in the pinv JVP.

2022-11-08 08:53:00 -05:00

..

_src

Add support for the hermitian option on jnp.linalg.pinv.

2022-11-08 08:53:00 -05:00

example_libraries

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

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

experimental

Fix GDA error message formatting

2022-11-07 11:55:28 -08:00

image

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

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

interpreters

Add UnloadedMeshExecutable to represent a MeshExecutable that is not loaded

2022-11-07 13:40:40 -08:00

lax

Add lax.cumlogsumexp for cumulative logsumexp operations.

2022-10-31 15:08:52 -07:00

lib

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

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

nn

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

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

numpy

Point to ndarray.at from docstring of unimplemented jnp.put & jnp.place

2022-10-28 14:13:36 -07:00

ops

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

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

scipy

implement truncnorm in jax.scipy.stats

2022-10-22 15:48:20 -04:00

tools

Prepare for release of jax and jaxlib 0.3.24

2022-11-03 15:13:23 -07:00

__init__.py

Add a new jax.spmd_mode config for preventing unintentional hangs and incorrect results when users pass jax.Arrays that span across multiple processes (i.e. not fully addressable) to jit or jnp operations (that are jitted by default).

2022-10-31 09:51:42 -07:00

abstract_arrays.py

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

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

ad_checkpoint.py

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

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

api_util.py

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

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

BUILD

Move array.py and sharding.py from experimental/ to _src/.

2022-09-27 10:06:52 -07:00

cloud_tpu_init.py

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

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

collect_profile.py

Cleanup: remove unused imports in private modules

2022-10-20 14:37:21 -07:00

config.py

Cleanup: remove unused imports in private modules

2022-10-20 14:37:21 -07:00

core.py

Add a number of missing function cross-references in the docs.

2022-11-07 12:00:26 -05:00

custom_batching.py

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

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

custom_derivatives.py

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

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

custom_transpose.py

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

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

debug.py

Add inspect_array_sharding, enabling looking at shardings in pjit-ted functions

2022-09-22 17:36:56 -07:00

distributed.py

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

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

dlpack.py

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

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

dtypes.py

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

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

errors.py

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

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

flatten_util.py

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

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

jaxpr_util.py

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

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

linear_util.py

add a basic prototype of piles, behind jax_dynamic_shapes

2022-11-06 17:03:04 -08:00

prng.py

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

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

profiler.py

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

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

py.typed

Improve type annotations

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

random.py

Add a number of missing function cross-references in the docs.

2022-11-07 12:00:26 -05:00

sharding.py

Rename ReshapeableDevicesSharding to PositionalSharding and add an alias NamedSharding for MeshPspecSharding.

2022-11-02 19:13:13 -07:00

stages.py

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

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

test_util.py

Remove deprecated functionality from jax.test_util

2022-10-11 08:16:34 -07:00

tree_util.py

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

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

util.py

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

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

version.py

Finish the release of jax and jaxlib 0.3.24

2022-11-04 09:43:12 -07:00