jax authors 0a84db59ed Merge pull request #6068 from jakevdp:fix-result-type ...

PiperOrigin-RevId: 363282198

2021-03-16 15:20:40 -07:00

..

_src

Merge pull request #6068 from jakevdp:fix-result-type

2021-03-16 15:20:40 -07:00

experimental

Add new lax.rng_bit_generator primitive

2021-03-16 16:30:09 +00:00

image

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

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

interpreters

Merge pull request #6048 from skye:debug_nans

2021-03-15 18:24:44 -07:00

lax

Add new lax.rng_bit_generator primitive

2021-03-16 16:30:09 +00:00

lib

Add extras_require to setup.py

2021-03-16 13:23:46 -07:00

nn

Delete jax.nn.functions.

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

numpy

Add test that docstring wrapping didn't fail

2021-03-04 14:16:43 -08: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

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

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

abstract_arrays.py

Fixed mypy type errors for numpy 1.20

2021-02-05 10:40:47 +02: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

Increase minimum jaxlib version to 0.1.62.

2021-03-16 15:11:36 -04: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

Increase minimum jaxlib version to 0.1.62.

2021-03-16 15:11:36 -04:00

core.py

Merge pull request #5977 from apaszke:xmap-with-control-flow

2021-03-09 11:22:18 -08:00

custom_derivatives.py

document what closure_convert returns

2021-03-02 09:18:46 -08: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

Merge pull request #6068 from jakevdp:fix-result-type

2021-03-16 15:20:40 -07:00

errors.py

Add jax.errors submodule & error troubleshooting docs

2021-03-03 12:39:12 -08:00

flatten_util.py

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

2021-01-11 14:21:07 -08: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

Fix lax_reference population_count implementation to avoid overflowing Python int -> np.uint32 casts.

2020-11-20 14:50:10 -05:00

lazy.py

Remove the JAX lazy sublanguage.

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

linear_util.py

Make jax_enable_x64 a thread-local value.

2021-02-04 09:48:22 -08: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

Merge pull request #5998 from zhangqiaorjc:dev_put_count

2021-03-10 14:36:55 -08: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

update jax for pypi

2021-03-05 10:12:58 -08:00