Yash Katariya 13c34f9dc5 Move with_sharding_constraint out of experimental into jax.lax namespace. ...

PiperOrigin-RevId: 494635809

2022-12-11 22:55:21 -08:00

..

_src

Improves handling of opaque types for dynamic shapes

2022-12-12 05:19:04 +01:00

example_libraries

Revert part of #13498

2022-12-05 09:21:09 -08:00

experimental

Improves handling of opaque types for dynamic shapes

2022-12-12 05:19:04 +01:00

image

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

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

interpreters

Improves handling of opaque types for dynamic shapes

2022-12-12 05:19:04 +01:00

lax

Move with_sharding_constraint out of experimental into jax.lax namespace.

2022-12-11 22:55:21 -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 jax.scipy.linalg.toeplitz.

2022-12-09 01:03:21 +09:00

tools

Prepare for jax and jaxlib 0.4.0 release

2022-12-07 16:02:24 -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

Add CUDNN custom call for LSTM. Exposed as jax.experimental.rnn module.

2022-11-28 14:31:48 -08: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

[jax2tf] An alternative support for shape polymorphism for native serialization.

2022-12-08 08:19:35 +02: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

monitoring.py

Add JAX monitoring library that instruments code via events.

2022-11-15 12:41:41 -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

Use jax.Array by default for doctests

2022-11-15 11:52:22 -08:00

sharding.py

Move PartitionSpec and Mesh out of experimental and into the sharding namespace. The new API endpoint is jax.sharding.PartitionSpec and jax.sharding.Mesh.

2022-12-01 19:28:32 -08: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 jax and jaxlib release 0.3.25. The next release will be 0.4.0 (since jax.Array will be enabled in that release)

2022-11-15 09:02:53 -08:00