Yash Katariya a683186570 Use the 11/09 libtpu build for jaxlib release since that passes all the tests. ...

PiperOrigin-RevId: 488543322

2022-11-14 20:37:41 -08:00

..

_src

Remove custom potrf kernels in favor of native XLA cholesky support.

2022-11-14 18:45:25 -08:00

example_libraries

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

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

experimental

Make MeshPspecSharding an alias for NamedSharding (it was the other way around before this CL).

2022-11-14 14:44:00 -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 docstrings for Sharding classes. Right now I am only documenting Sharding, XLACompatibleSharding, MeshPspecSharding and SingleDeviceSharding.

2022-11-14 15:47:46 -08:00

lax

Add support for Hessenberg and tridiagonal matrix reductions on CPU.

2022-11-09 06:23:55 -08: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

Add support for Hessenberg and tridiagonal matrix reductions on CPU.

2022-11-09 06:23:55 -08:00

tools

Use the 11/09 libtpu build for jaxlib release since that passes all the tests.

2022-11-14 20:37:41 -08: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

Delete jax_experimental_name_stack flag

2022-11-10 11:59:50 -08: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