Peter Hawkins 0150d15cb2 Increase minimum jaxlib version to 0.3.7. ...

Drop backwards compatibility with older jaxlib versions.

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

..

compilation_cache

Increase minimum jaxlib version to 0.3.7.

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

gda_serialization

Enable use of GlobalDeviceArray (GDA) in T5X Checkpointer. Add a separate unit test, gda_checkpoints_test, to cover this use case.

2022-04-04 10:56:07 -07:00

jax2tf

Merge pull request #9809 from pschuh:opt-barrier-gpu

2022-04-15 12:43:50 -07:00

sparse

Increase minimum jaxlib version to 0.3.7.

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

__init__.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

callback.py

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

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

checkify.py

Restructure checkify files.

2022-04-14 04:32:24 -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

Adds simple effect types to jaxprs

2022-04-11 11:50:41 -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

Add tests for higher order primitives

2022-04-12 18:12:44 -07:00

mesh_utils.py

Remove unused comments related to Python 2 compatibility.

2022-04-14 12:52:51 -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

Drop dead code now that the minimum jaxlib version is 0.3.2.

2022-04-13 13:34:00 -04: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