Yash Katariya 1ad3551ec9 Release jax and jaxlib 0.3.0 as per the new release process. ...

PiperOrigin-RevId: 427809845

2022-02-10 11:59:13 -08:00

..

_src

Add complex number DLPack support to JAX and TensorFlow.

2022-02-09 14:58:00 -08:00

example_libraries

[JAX] move example libraries from jax.experimental into jax.example_libraries

2021-10-19 17:30:45 -07:00

experimental

Checkify: Fix docstring formatting and polish enabled_errors sets.

2022-02-10 16:54:43 +00:00

image

Add PEP484-compatible export for jax and its subpackages

2021-09-13 14:08:48 -07:00

interpreters

Strip named_shape and weak_type from aval when donating buffers.

2022-02-10 07:39:50 -08:00

lax

Change lax.select_p to be an n-ary predicate, 'lax.select_n_p'. Change lax.select() to be a thin shim around the new n-ary version.

2022-02-09 11:03:09 -08:00

lib

Use imports relative to the jax package consistently, rather than .-relative imports.

2021-11-24 07:48:29 -08:00

nn

Use imports relative to the jax package consistently, rather than .-relative imports.

2021-11-24 07:48:29 -08:00

numpy

Add jax.numpy.generic

2022-02-07 14:56:39 -08:00

ops

Add PEP484-compatible export for jax and its subpackages

2021-09-13 14:08:48 -07:00

scipy

Merge pull request #8625 from Edenhofer:regular_grid_interpolator

2022-01-24 12:02:34 -08:00

tools

Fix SyntaxError in jax_to_ir

2021-12-21 15:11:08 -08:00

__init__.py

Rollback of: Implement the JAX transfer guard API

2022-02-09 14:44:45 -08:00

abstract_arrays.py

Move jax.abstract_arrays to jax._src.abstract_arrays.

2021-06-02 06:25:22 -07:00

ad_checkpoint.py

checkpoint_name for checkpoint policies by name

2021-10-13 21:31:02 -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

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

Replace core.pp_eqn_compact() with core.str_eqn_compact().

2022-02-09 22:08:45 +00:00

custom_batching.py

introduce custom_batching.sequential_vmap

2022-01-20 21:37:58 -08:00

custom_derivatives.py

Add PEP484-compatible export for jax and its subpackages

2021-09-13 14:08:48 -07:00

custom_transpose.py

introduce custom_transpose and a corresponding primitive

2022-01-11 12:51:17 -08:00

distributed.py

Add jax.distributed.initialize for multi-host GPU.

2021-10-26 14:37:54 -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

Add PEP484-compatible export for jax and its subpackages

2021-09-13 14:08:48 -07:00

errors.py

Use imports relative to the jax package consistently, rather than .-relative imports.

2021-11-24 07:48:29 -08: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

Speed up source location computation when lowering a jaxpr to HLO/MHLO.

2022-02-08 16:17:40 -08:00

linear_util.py

clean up WrappedFun.call_wrapped refs on exception

2022-01-05 16:11:48 -08:00

prng.py

improvements to RBG PRNG

2021-10-07 18:59:13 -07:00

profiler.py

Add PEP484-compatible export for jax and its subpackages

2021-09-13 14:08:48 -07:00

py.typed

Improve type annotations

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

random.py

add random.default_prng_impl to retrieve the default PRNG implementation

2022-01-12 19:13:14 -08:00

test_util.py

test_util: export with_config

2022-01-26 14:32:11 -08:00

tree_util.py

Add PEP484-compatible export for jax and its subpackages

2021-09-13 14:08:48 -07:00

util.py

Remove jax.util.partial.

2021-09-28 11:04:02 -04:00

version.py

Release jax and jaxlib 0.3.0 as per the new release process.

2022-02-10 11:59:13 -08:00