Skye Wanderman-Milne 0a69c9a27b Fix jax2tf import so it works with both the latest tensorflow release (2.10.0) and tf-nightly

2022-09-30 15:55:22 -07:00

..

_src

Merge pull request #12594 from skye:cache_warnings

2022-09-30 13:11:34 -07: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 jax2tf import so it works with both the latest tensorflow release (2.10.0) and tf-nightly

2022-09-30 15:55:22 -07:00

image

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

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

interpreters

Disable dynamic_api_test and custom_object_test.py with jax.Array. Enable it back when support for it is added. Also don't use xla_shape since its deprecated.

2022-09-29 15:09:55 -07:00

lax

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

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

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

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

ops

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

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

scipy

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

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

tools

Pin default jax.tools.colab_tpu.setup_tpu driver version.

2022-09-30 17:45:49 +00:00

__init__.py

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

2022-09-27 10:06:52 -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

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

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

config.py

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

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

core.py

[dynamic-shapes] Add basic slicing support

2022-09-28 15:55:51 -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

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

fix bug in djax type signature inference logic

2022-09-26 17:48:25 -07: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

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

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

sharding.py

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

2022-09-27 10:06:52 -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

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

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

Update version numbers after release.

2022-09-28 18:49:22 +00:00