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

..

clusters

Migrate JAX internals to builtin Python logging

2022-10-13 21:32:44 +02:00

debugger

Cleanup: remove unused imports in private modules

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

image

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

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

lax

Merge pull request #13147 from hawkinsp:eyes

2022-11-07 19:25:15 -08:00

lib

Increase the minimum jaxlib version to 0.3.22.

2022-10-27 10:24:11 -04:00

nn

Avoid loading scipy eagerly.

2022-10-12 19:51:09 +00:00

numpy

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

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

ops

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

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

scipy

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

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

state

Enable state effect in cond_p (except in grad and vmap)

2022-11-02 16:07:58 -07:00

third_party

Call _check_arraylike for jnp.linalg & jnp.fft functions

2022-10-31 09:19:53 -07:00

__init__.py

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

2022-09-22 12:27:19 -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

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

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

ad_util.py

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

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

api_util.py

Annotate tree_util

2022-10-26 13:38:38 -04:00

api.py

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

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

array.py

Add global_shards to jax.Array as it exists on GDA and is being used in various places.

2022-10-31 09:08:03 -07:00

basearray.py

Add initial jax.Array base class for instance checks & annotation

2022-09-26 07:48:43 -07:00

basearray.pyi

[typing] overloads for jnp.linalg.svd & jnp.linalg.qr

2022-10-05 14:29:00 -07:00

callback.py

Increase the minimum jaxlib version to 0.3.22.

2022-10-27 10:24:11 -04:00

checkify.py

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

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

cloud_tpu_init.py

Set JAX_PLATFORMS=tpu,cpu on TPUs

2022-10-05 04:39:04 +00:00

config.py

Merge pull request #13039 from skye:cache_compile_time_heuristic

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

custom_api_util.py

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

2022-09-22 12:27:19 -07: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

Cleanup: remove unused imports in private modules

2022-10-20 14:37:21 -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

debugging.py

Merge pull request #13065 from sharadmv:vis-colors

2022-11-01 18:11:06 -07:00

device_array.py

Add addressable_shards to SDA and DA as a compatibility API to match with jax.Array. This will aid in transition to jax.Array.

2022-10-05 12:35:17 -07:00

dispatch.py

Improve the error message when the avals a function was AOT compiled with doesn't match the input avals when its called.

2022-11-04 21:25:46 -07:00

distributed.py

Migrate JAX internals to builtin Python logging

2022-10-13 21:32:44 +02:00

dlpack.py

canonicalize dtypes when loading arrays via dlpack

2022-10-20 15:20:34 -07:00

dtypes.py

[typing] annotate jax._src.dtypes

2022-10-17 13:49:26 -07:00

environment_info.py

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

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

errors.py

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

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

flatten_util.py

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

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

iree.py

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

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

lax_reference.py

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

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

lazy_loader.py

Avoid loading scipy eagerly.

2022-10-12 19:51:09 +00:00

path.py

Make the etils[epath] dependency of JAX optional.

2022-10-12 21:16:48 +00:00

pretty_printer.py

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

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

prng.py

integrate new (partitionable, count-space-exhaustive) counts generation

2022-10-29 00:05:49 -07:00

profiler.py

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

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

public_test_util.py

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

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

random.py

fix jax.random.logits shape argument

2022-11-04 19:51:39 -07:00

sharding.py

Remove device_indices method which is redundant because of the existence of devices_indices_map and is slower because pinging a cache for every device is not free.

2022-11-05 17:33:37 -07:00

source_info_util.py

Include column information in Python locations under Python 3.11.

2022-10-14 19:14:35 +00:00

stages.py

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

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

test_util.py

Remove error-prone most_recent_entry() support from lu.cache.

2022-10-27 16:41:44 -07:00

traceback_util.py

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

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

tree_util.py

Annotate tree_util

2022-10-26 13:38:38 -04:00

typing.py

add allow_opaque_dtype to dtypes.canonicalize_dtype utility

2022-10-17 13:47:42 -07:00

util.py

Increase the minimum jaxlib version to 0.3.22.

2022-10-27 10:24:11 -04:00