Peter Hawkins fa411d864e [MLIR] Fix CPU test failures for MLIR lowering. ...

The remaining failures relate to buffer donation and xmap_p, which are not yet implemented.

Quite a few primitives still use fallback paths.

PiperOrigin-RevId: 413130158

2021-11-30 06:08:55 -08:00

..

compilation_cache

Increase minimum jaxlib version to 0.1.74.

2021-11-18 15:06:58 -05:00

jax2tf

[jax2tf] Implement CLIP mode for scatter.

2021-11-26 10:01:42 +02:00

sparse

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

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

[MLIR] Fix CPU test failures for MLIR lowering.

2021-11-30 06:08:55 -08:00

callback.py

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

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

djax.py

Merge xla._partition_outputs and util.unflatten.

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

global_device_array.py

Rename GlobalShardedDeviceArray to GlobalDeviceArray

2021-11-24 16:55:29 -08:00

host_callback.py

Remove type annotations for id_tap.

2021-11-29 14:25:10 -08:00

jet.py

Split jax.interpreters.xla up into three pieces:

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

loops.py

Switch internal users of jax.ops.index_... to use x.at[x].set() APIs.

2021-09-13 19:48:29 -04:00

maps.py

ahead-of-time lowering and compilation frontend for xmap

2021-11-24 10:49:24 -08:00

ode.py

odeint: args validation allows pytrees

2021-10-21 16:34:59 -07:00

optimizers.py

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

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

pjit.py

Rename GlobalShardedDeviceArray to GlobalDeviceArray

2021-11-24 16:55:29 -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

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

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