Lena Martens 7b5b9cefbd Add scatter OOB error.

2022-01-07 17:22:34 +00: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

Merge pull request #8974 from gnecula:refactor_poly

2021-12-16 02:44:56 -08:00

sparse

[sparse] add getitem stub in JAXSparse

2022-01-05 12:42:30 -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 scatter OOB error.

2022-01-07 17:22:34 +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

Replace TODO with a different one

2021-12-16 17:58:21 -08:00

host_callback.py

Merge pull request #8934 from juesato:hcb_docs

2021-12-15 08:05:32 -08: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

[MLIR] Add a @mlir.cache_lowering decorator that lowers a primitive out-of-line as a reusable function.

2021-12-16 08:34:52 -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 comment clarifying why gda_ppspec is returned.

2021-12-16 05:44:16 -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