Matthew Johnson 70cd62d545 disable disable_omnistaging

2021-03-23 19:13:15 -07:00

..

_src

unify configuration state handling

2021-03-23 18:56:01 -07:00

experimental

unify configuration state handling

2021-03-23 18:56:01 -07:00

image

Move implementation of jax.image to jax._src.image.

2020-10-17 11:37:55 -04:00

interpreters

unify configuration state handling

2021-03-23 18:56:01 -07:00

lax

add clz primitive

2021-03-19 22:54:36 -07:00

lib

Specify minimum jaxlib version in a single location

2021-03-22 16:14:41 -07:00

nn

Delete jax.nn.functions.

2020-11-09 19:48:54 -05:00

numpy

Initial implementation of jnp.delete

2021-03-16 17:05:23 -07:00

ops

Move jax.ops implementation into jax._src.ops.

2020-10-17 11:45:28 -04:00

scipy

added bicgstab to new jax repo

2021-02-18 18:01:28 -08:00

tools

Merge pull request #5173 from jakevdp:conv-elem-type-cleanup

2020-12-11 15:27:40 -08:00

__init__.py

unify configuration state handling

2021-03-23 18:56:01 -07:00

abstract_arrays.py

fix convert_element_type on large Py int inputs

2021-03-21 19:08:59 -07:00

ad_util.py

Some minor typing fixes

2021-03-02 09:42:52 -08:00

api_util.py

improve error messages for grad(..., has_aux=True)

2021-02-18 09:46:16 -08:00

api.py

unify configuration state handling

2021-03-23 18:56:01 -07:00

BUILD

Merge pull request #5717 from google:dynamic-shapes2

2021-02-16 18:45:13 -08:00

cloud_tpu_init.py

Automatically initialize Cloud TPU topology env vars if running on a Cloud TPU VM.

2021-03-10 09:15:31 -08:00

config.py

disable disable_omnistaging

2021-03-23 19:13:15 -07:00

core.py

unify configuration state handling

2021-03-23 18:56:01 -07:00

custom_derivatives.py

unify configuration state handling

2021-03-23 18:56:01 -07:00

dlpack.py

[call_tf] Add support for DLPack to avoid copying arrays in eager mode

2021-02-13 12:49:51 +02:00

dtypes.py

unify configuration state handling

2021-03-23 18:56:01 -07:00

errors.py

errors: add NonConcreteBooleanIndexError & debugging tips

2021-03-23 11:23:20 -07:00

flatten_util.py

generalize ravel_pytree to handle int types, add tests

2021-03-19 10:50:02 -07:00

jax.bzl

Move jax.bzl into correct location. Unbreaks build.

2021-01-21 06:56:23 -08:00

jaxpr_util.py

[JAX] Move jax.util to jax._src_util.

2021-01-11 14:21:07 -08:00

lax_reference.py

add clz primitive

2021-03-19 22:54:36 -07:00

lazy.py

Remove the JAX lazy sublanguage.

2021-03-09 21:40:46 -05:00

linear_util.py

unify configuration state handling

2021-03-23 18:56:01 -07:00

profiler.py

[JAX] Move implementation of profiler.py into jax._src.

2020-11-04 09:06:40 -08:00

py.typed

Improve type annotations

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

random.py

DOC: fix jax.random documentation

2020-11-20 16:41:32 -08:00

test_util.py

unify configuration state handling

2021-03-23 18:56:01 -07:00

tree_util.py

Move flake8 & mypy checks to pre-commit

2021-02-10 13:44:40 -08:00

util.py

Export split_dict from jax.util

2021-01-29 11:13:36 -08:00

version.py

Specify minimum jaxlib version in a single location

2021-03-22 16:14:41 -07:00