Marc van Zee 9b54660183 [jax2tf] Use tf.nest.flatten instead of tf.keras.backend.flatten. ...

PiperOrigin-RevId: 460395701

2022-07-12 00:54:06 -07:00

..

_src

Define PyTreeDef on pytree module rather than its parent

2022-07-11 17:57:36 -07:00

example_libraries

feat: refactor code using pyupgrade

2022-05-17 22:14:05 +01:00

experimental

[jax2tf] Use tf.nest.flatten instead of tf.keras.backend.flatten.

2022-07-12 00:54:06 -07:00

image

CI: fix flake8 ignore declarations

2022-04-21 13:44:12 -07:00

interpreters

Merge pull request #11405 from mattjj:djax-vmap

2022-07-09 10:38:39 -07:00

lax

Add bitwise XOR reducer to lax.reduce

2022-06-15 16:56:51 +02:00

lib

Reexport jaxlib.__version as jax.lib.__version__.

2022-04-28 10:25:06 -07:00

nn

CI: fix flake8 ignore declarations

2022-04-21 13:44:12 -07:00

numpy

CI: fix flake8 ignore declarations

2022-04-21 13:44:12 -07:00

ops

CI: fix flake8 ignore declarations

2022-04-21 13:44:12 -07:00

scipy

Implementation of jax.scipy.stats.gaussian_kde

2022-06-28 15:17:12 -04:00

tools

Add support for running JAX tests under Bazel.

2022-07-01 15:07:22 -07:00

__init__.py

Deprecate top-level aliases of jax.tree_util functions

2022-07-07 11:41:46 -07:00

abstract_arrays.py

CI: fix flake8 ignore declarations

2022-04-21 13:44:12 -07:00

ad_checkpoint.py

CI: fix flake8 ignore declarations

2022-04-21 13:44:12 -07:00

api_util.py

CI: fix flake8 ignore declarations

2022-04-21 13:44:12 -07:00

BUILD

Merge pull request #11359 from hawkinsp:bazel

2022-07-06 06:13:20 -07:00

cloud_tpu_init.py

CI: fix flake8 ignore declarations

2022-04-21 13:44:12 -07:00

collect_profile.py

Add collect_profile script

2022-06-03 17:56:17 -07:00

config.py

CI: fix flake8 ignore declarations

2022-04-21 13:44:12 -07:00

core.py

Merge pull request #11405 from mattjj:djax-vmap

2022-07-09 10:38:39 -07:00

custom_batching.py

CI: fix flake8 ignore declarations

2022-04-21 13:44:12 -07:00

custom_derivatives.py

CI: fix flake8 ignore declarations

2022-04-21 13:44:12 -07:00

custom_transpose.py

CI: fix flake8 ignore declarations

2022-04-21 13:44:12 -07:00

distributed.py

Unbreak jax.distributed initialization.

2022-06-07 11:10:17 -07:00

dlpack.py

CI: fix flake8 ignore declarations

2022-04-21 13:44:12 -07:00

dtypes.py

CI: fix flake8 ignore declarations

2022-04-21 13:44:12 -07:00

errors.py

CI: fix flake8 ignore declarations

2022-04-21 13:44:12 -07:00

flatten_util.py

CI: fix flake8 ignore declarations

2022-04-21 13:44:12 -07:00

jaxpr_util.py

Fix or ignore some pytype errors.

2022-05-31 21:25:55 -07:00

linear_util.py

[dynamic-shapes] basic linearize and grad working

2022-06-30 14:30:22 -07:00

prng.py

CI: fix flake8 ignore declarations

2022-04-21 13:44:12 -07:00

profiler.py

Store profiler server as a global variable and add a stop_server function

2022-06-01 17:50:06 -07:00

py.typed

Improve type annotations

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

random.py

Fix RST formatting in random.py docstring

2022-07-11 02:51:35 -07:00

stages.py

CI: fix flake8 ignore declarations

2022-04-21 13:44:12 -07:00

test_util.py

jax.test_util: remove deprecated test classes.

2022-06-27 11:04:50 -07:00

tree_util.py

CI: fix flake8 ignore declarations

2022-04-21 13:44:12 -07:00

util.py

CI: fix flake8 ignore declarations

2022-04-21 13:44:12 -07:00

version.py

Update minimum jaxlib version to 0.3.14.

2022-07-08 00:36:02 +00:00