jax authors 8e6101c6a1 Merge pull request #6866 from gnecula:tf_pjit ...

PiperOrigin-RevId: 376989780

2021-06-01 22:50:12 -07:00

..

_src

Merge pull request #6785 from GregCT:changelist/373551581

2021-06-01 14:47:15 -07:00

experimental

Merge pull request #6866 from gnecula:tf_pjit

2021-06-01 22:50:12 -07:00

image

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

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

interpreters

[JAX] Change how JAX manages XLA platforms.

2021-06-01 11:44:31 -07:00

lax

Add lax.reduce_precision()

2021-04-05 09:54:14 -07:00

lib

[JAX] Change how JAX manages XLA platforms.

2021-06-01 11:44:31 -07:00

nn

add tanh to jax.nn package

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

numpy

Add jnp.resize

2021-05-21 09:25:49 +01:00

ops

Implement segment_prod, segment_max, segment_min

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

scipy

Add a tridiagonal eigh solver.

2021-05-04 20:43:41 -04:00

tools

[host_callback] Minor fix to use the new xla_shape.is_token

2021-04-28 12:22:32 +03:00

__init__.py

Replace host_id with process_index terminology, take 2.

2021-04-20 18:13:34 -07:00

abstract_arrays.py

[jax2tf] Implementation of a parametric shape-polymorphism feature for jax2tf.

2021-04-08 10:42:38 +03:00

ad_util.py

Some minor typing fixes

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

api_util.py

improve concreteness error from arguments

2021-05-03 17:37:34 -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

handle case where trace debug_info is None

2021-05-06 18:38:20 -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

errors: add NonConcreteBooleanIndexError & debugging tips

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

flatten_util.py

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

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

jaxpr_util.py

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

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

linear_util.py

Handle context manager configuration settings for matmul precision and numpy rank promotion correctly in JIT and linear_util caches.

2021-04-21 06:36:35 -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

Run doctest on all source files except jax2tf

2021-04-05 10:39:59 -07:00

test_util.py

[jax2tf] Add support for pjit.

2021-06-01 14:32:59 +03:00

tree_util.py

Move jax.tree_util implementation to jax._src.tree_util.

2021-03-24 12:00:38 -07:00

util.py

Export split_dict from jax.util

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

version.py

Bump jax version and changelog for jax 0.1.13 release

2021-05-03 16:32:00 -07:00