George Necula f08156ab7c [hcb] Simplifications to the host_calback API ...

* dropping support for special AD handling for hcb.id_tap and id_print.
  From now on, only the primals are tapped. The old behavior can be
  obtained (for a limited time) by setting the JAX_HOST_CALLBACK_AD_TRANSFORMS
  environment variale, or the --flax_host_callback_ad_transforms flag.
  Additionally, added documentation for how to implement the old behavior
  using JAX custom AD APIs.

This allows us to make some significant cleanup in the internals.

2021-12-11 08:24:56 +01:00

..

image

Added workaround for bug in XLA

2021-11-18 11:01:50 +02:00

lax

[MLIR] Support all fill_modes in GPU MLIR lowering for scatter_add.

2021-12-10 15:01:08 -08:00

lib

Consolidate more XLA-lowering logic between jit, pmap, and xmap.

2021-11-30 14:24:33 -08:00

nn

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

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

numpy

[x64] minor weak_type changes to linalg.py

2021-12-07 16:27:29 -08:00

ops

[x64] ensure scatter functionality preserves weak_type

2021-11-30 15:43:06 -08:00

scipy

[x64] minor weak_type changes to linalg.py

2021-12-07 16:27:29 -08: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

Move dtype canonicalization out of core.AbstractValue subclasses.

2021-12-07 06:13:07 -08:00

ad_checkpoint.py

Copybara import of the project:

2021-12-01 08:43:09 -08:00

ad_util.py

Move jax.ad_util to jax._src.ad_util.

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

api_util.py

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

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

api.py

[JAX] Add a dialect option to jit(...).lower(...).compiler_ir().

2021-12-10 13:02:25 -08: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

[hcb] Simplifications to the host_calback API

2021-12-11 08:24:56 +01:00

custom_derivatives.py

[MLIR] Move most MLIR translation rules into lax.

2021-11-23 18:58:28 -08:00

device_array.py

Move dtype canonicalization out of core.AbstractValue subclasses.

2021-12-07 06:13:07 -08:00

dispatch.py

[MLIR] Keep MLIR IR longer as a Python ir.Module object rather than a string, until it is time to compile it.

2021-12-10 14:56:48 -08:00

distributed.py

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

2021-10-26 14:37:54 -07:00

dlpack.py

Split jax.interpreters.xla up into three pieces:

2021-11-22 08:22:43 -08:00

dtypes.py

[x64] add new jax_default_dtype_bits flag

2021-12-07 08:23:37 -08: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

iree.py

Add a prototype IREE backend for JAX.

2021-11-15 07:57:04 -08:00

lax_reference.py

Improve implementation of cbrt() in JAX.

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

pretty_printer.py

Adds a Wadler-Lindig pretty printer.

2021-09-27 21:09:24 -04:00

prng.py

Overload jnp ops which are polymorphic to an array's value and support PRNGKeys.

2021-11-16 23:00:32 +00:00

profiler.py

Add profiler annotations

2021-12-06 15:13:35 -08:00

random.py

Cleanup: remove duplicate canonicalize_axis utility

2021-11-23 16:54:02 -08:00

source_info_util.py

Fix forward for PR #8392 (made source_info for new_jaxpr_eqn argument optional again)

2021-10-29 15:50:14 -07:00

test_util.py

[hcb] Simplifications to the host_calback API

2021-12-11 08:24:56 +01:00

traceback_util.py

Move contents of jax.lib to jax._src.lib.

2021-09-23 06:33:55 -07:00

tree_util.py

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

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

util.py

Merge xla._partition_outputs and util.unflatten.

2021-11-24 12:52:40 -08:00