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

..

_src

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

experimental

New lowering APIs for pjit

2021-09-08 10:54:41 -07:00

image

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

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

interpreters

New lowering APIs for pjit

2021-09-08 10:54:41 -07:00

lax

ReduceScatter translation and abstract eval.

2021-07-27 11:13:18 -07:00

lib

Drop support for the deprecated StreamExecutor CPU backend.

2021-09-08 06:04:45 -07:00

nn

add tanh to jax.nn package

2021-04-29 08:26:38 -07:00

numpy

Implement np.polyfit

2021-08-31 10:52:45 +05:30

ops

Implement segment_prod, segment_max, segment_min

2021-04-09 12:06:51 -07:00

scipy

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

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

tools

Unpin the tpu_driver version used for Cloud TPU Colabs.

2021-09-07 16:12:36 -07:00

__init__.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

abstract_arrays.py

Move jax.abstract_arrays to jax._src.abstract_arrays.

2021-06-02 06:25:22 -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

Replace host_id with process_index terminology, take 2.

2021-04-20 18:13:34 -07: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

Move jax.config to jax._src.config.

2021-04-19 08:53:12 -07:00

core.py

Remove axis name from named_shape when unmapping avals

2021-09-08 01:42:15 -07:00

custom_derivatives.py

Move contents of jax.custom_derivatives to jax._src.custom_derivatives.

2021-04-19 17:51:49 -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

[JAX] Move contents of jax/dtypes.py to jax/_src/dtypes.py.

2021-04-07 19:35:51 -07:00

errors.py

Expose UnexpectedTracerError and add docs.

2021-07-27 23:23:28 +01: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

jaxpr_util.py

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

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

linear_util.py

Cleanup: switch to new version of super()

2021-08-05 13:11:07 -07:00

prng.py

implement backwards-compatible behavior and enable custom PRNGs only conditionally

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

profiler.py

Add programmatic profiling APIs, and rename some existing APIs.

2021-03-25 11:50:55 -07:00

py.typed

Improve type annotations

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

random.py

alias prng.threefry_2x32 in random and warn of move

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

test_util.py

fix custom_vjp issue?

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

tree_util.py

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

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

util.py

Export split_dict from jax.util

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

version.py

update version and changelog for pypi

2021-09-02 15:38:47 -07:00