George Necula 3021d3e2e2 [hcb] Add support for remat2 to host_callback ...

A callback under ad_checkpoint.checkpoint will be invoked
twice when taking the gradient: once during the forward pass
and once again during the backward pass when the residuals
for the forward pass are rematerialized.

2021-12-15 10:32:15 +02: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

Adds rtol to jax2tf/converters_eval tests.

2021-12-14 03:32:32 -08:00

sparse

[sparse] add sparse rule for lax.sub_p

2021-12-13 16:46:47 -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

Add a fast_path to GDA for faster GDA construction. Local indices

2021-12-14 18:53:56 -08:00

host_callback.py

[hcb] Add support for remat2 to host_callback

2021-12-15 10:32:15 +02: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

log tracing and compilation times

2021-12-13 21:58:09 -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

Add a fast_path to GDA for faster GDA construction. Local indices

2021-12-14 18:53:56 -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