Sharad Vikram 098f2126ae Add CustomCall MLIR lowering for HCB outside_call primitive ...

PiperOrigin-RevId: 444436860

2022-04-25 19:38:39 -07:00

..

compilation_cache

Increase minimum jaxlib version to 0.3.7.

2022-04-18 08:09:50 -04:00

gda_serialization

Add a dtypes option to cast host arrays when reloading from TS.

2022-04-22 18:00:27 -07:00

jax2tf

CI: fix flake8 ignore declarations

2022-04-21 13:44:12 -07:00

sparse

Merge pull request #10443 from jakevdp:sparse-notimplemented

2022-04-25 19:03:02 -07:00

__init__.py

CI: fix flake8 ignore declarations

2022-04-21 13:44:12 -07:00

callback.py

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

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

checkify.py

CI: fix flake8 ignore declarations

2022-04-21 13:44:12 -07:00

global_device_array.py

Move PartitionSpec from sharded_jit.py to pxla.py. The public endpoint is via jax.experimental so that should be used (no changes to the public endpoint).

2022-04-06 15:19:19 -07:00

host_callback.py

Add CustomCall MLIR lowering for HCB outside_call primitive

2022-04-25 19:38:39 -07:00

jet.py

remove _reduce_min from public jax.lax module

2022-03-08 16:34:26 -08:00

loops.py

Deprecate jax.experimental.loops, step 2.

2022-04-14 12:38:55 -07:00

maps.py

Merge pull request #10402 from froystig:aot-jit-avoid-trivial

2022-04-21 18:13:10 -07:00

mesh_utils.py

[mesh_utils] Add device/slice count checks

2022-04-20 13:30:01 -07:00

multihost_utils.py

Move PartitionSpec from sharded_jit.py to pxla.py. The public endpoint is via jax.experimental so that should be used (no changes to the public endpoint).

2022-04-06 15:19:19 -07:00

ode.py

Fix a bug that promoted t to a complex in odeint, and modify a test so it would have caught it

2022-01-07 12:10:35 -05: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 in mlir lowering for tokens

2022-04-21 11:28:58 -07: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