Peter Hawkins e7e5140dc9 Move implementation of jax.flatten_util to jax._src.flatten_util. Add a jax.flatten_util shim. ...

Change as part of cleaning up the jax.* namespace.

PiperOrigin-RevId: 395551093

2021-09-08 13:54:25 -07:00

..

image

jax.image: fix error with zero-size image resizing

2021-08-11 11:20:20 -07:00

lax

Merge pull request #7854 from hawkinsp:rename

2021-09-08 12:03:19 -07:00

nn

Merge branch 'main' into complex_init

2021-08-23 13:52:43 +02:00

numpy

Merge pull request #7843 from hawkinsp:linspace

2021-09-08 11:09:50 -07:00

ops

Validate shapes for boolean indices

2021-08-03 15:33:44 -07:00

scipy

add support for scipy.special.{expn,expi,exp1}

2021-08-24 16:36:10 -04:00

third_party

Move jax.third_party to jax._src.third_party.

2020-10-17 11:43:07 -07:00

__init__.py

Move most of the implementation of jax.scipy into jax._src.scipy.

2020-10-16 17:04:25 -04:00

abstract_arrays.py

Ensure zeros from AD are generated on device.

2021-07-26 20:40:13 +03:00

ad_util.py

Move jax.ad_util to jax._src.ad_util.

2021-06-07 14:51:34 -07:00

api_util.py

Move contents of jax.api_util to jax._src.api_util and add a forwarding shim.

2021-09-08 09:00:56 -07:00

api.py

Move contents of jax.api_util to jax._src.api_util and add a forwarding shim.

2021-09-08 09:00:56 -07:00

cloud_tpu_init.py

Changes to make jax[tpu] work better in a docker container.

2021-07-12 17:42:46 -07:00

config.py

implement backwards-compatible behavior and enable custom PRNGs only conditionally

2021-08-19 20:43:11 -07:00

custom_derivatives.py

Move contents of jax.api_util to jax._src.api_util and add a forwarding shim.

2021-09-08 09:00:56 -07:00

dlpack.py

Update minimum jaxlib version to 0.1.69.

2021-07-15 17:00:13 -04:00

dtypes.py

Fix dtypes.issubdtype when called with "bfloat16" (as string)

2021-05-23 19:32:45 +03:00

errors.py

Specify weak_type in DeviceArray repr

2021-08-23 13:19:33 -07:00

flatten_util.py

Move implementation of jax.flatten_util to jax._src.flatten_util. Add a jax.flatten_util shim.

2021-09-08 13:54:25 -07:00

lax_reference.py

Improve implementation of cbrt() in JAX.

2021-07-22 14:01:28 -07:00

pprint_util.py

[JAX] Move pprint_util into jax._src.

2020-11-19 06:42:19 -08:00

prng.py

inline jit-decorated jax.random calls

2021-08-20 13:43:38 -07:00

profiler.py

Add missing typing.Optional type annotations to function parameters.

2021-05-27 20:10:23 -07:00

random.py

inline jit-decorated jax.random calls

2021-08-20 13:43:38 -07:00

source_info_util.py

Most of JAX now uses concrete types for things defined in jaxlib.xla_client

2021-08-16 20:33:36 +01:00

traceback_util.py

fix custom_vjp issue?

2021-08-27 21:22:25 -07:00

tree_util.py

Specify weak_type in DeviceArray repr

2021-08-23 13:19:33 -07:00

util.py

address comments form @apaszke

2021-08-26 14:10:58 -07:00