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

..

compilation_cache

[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

gda_serialization

Add gcs support to serialization/deserialization of GDA

2021-12-09 14:25:26 -08:00

jax2tf

Testing: avoid hard-coding random seeds

2021-12-10 10:32:09 -08:00

sparse

[sparse] fix issue with duplicate indices in BCOO.__mul__

2021-12-10 09:37:35 -08:00

__init__.py

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

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

ann.py

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

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

callback.py

Switch internal users of jax.util.partial to use functools.partial.

2021-09-13 21:09:58 -04:00

checkify.py

Add checkify rule for while_loop.

2021-12-08 19:34:20 +00:00

djax.py

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

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

global_device_array.py

Docstring for GlobalDeviceArray and Shard.

2021-12-09 15:16:41 -08:00

host_callback.py

[hcb] Simplifications to the host_calback API

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

jet.py

Split jax.interpreters.xla up into three pieces:

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

loops.py

[x64] make jax.experimental.loops consistent with default dtype

2021-12-08 12:08:49 -08:00

maps.py

Thread is_gda to lower_mesh_computation. This is required because to lift the contiguous mesh requirement, the inputs should all be GDAs. The only way to check that in lower_mesh_computation when the input is being handled is to thread is_gda to lower_mesh_computation.

2021-12-06 15:44:03 -08:00

mesh_utils.py

Add contiguous_submeshes option to mesh_utils.create_device_mesh().

2021-12-10 00:01:12 +00:00

ode.py

Merge pull request #8339 from mariogeiger:main

2021-11-30 13:15:03 -08:00

optimizers.py

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

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

pjit.py

Clarify a comment in ParsedPartitionSpec.eq_given_rank

2021-12-10 11:30:15 -08:00

stax.py

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

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

x64_context.py

[x64] make x64_context tests work with updated default dtypes

2021-12-08 13:05:40 -08:00