George Necula cf4e1d414b [jax2tf] Bump the default JAX serialization version to 7. ...

This enables shape assertion checking, the support for which
landed in XlaCallModule on July 12th, 2023.

See the CHANGELOG for details.

PiperOrigin-RevId: 556222908

2023-08-11 22:49:31 -07:00

..

array_serialization

Add serialize_with_paths and deserialize_with_paths API to GlobalAsyncCheckpointManager

2023-08-08 22:47:27 -07:00

compilation_cache

Move implementation of compilation cache out of jax/experimental and into jax/_src.

2023-04-17 08:35:53 -07:00

jax2tf

[jax2tf] Bump the default JAX serialization version to 7.

2023-08-11 22:49:31 -07:00

mosaic

Add Mosaic to Jaxlib and expose bindings in jax.experimental.mosaic

2023-07-20 18:28:51 -07:00

pallas

[Pallas] Import ops in jax/experimental/pallas/ops/__init__.py

2023-08-01 22:41:30 -07:00

sparse

Add autodiff rules for spsolve_p

2023-08-04 12:08:57 -07:00

__init__.py

Copybara import of the project:

2023-05-24 10:35:37 -07:00

checkify.py

Document why 'import name as name' is used

2022-12-14 15:07:04 -08:00

custom_partitioning.py

Pad generated sharding specs with None up to ndims to simplify comparing dims

2023-08-10 17:02:31 -07:00

host_callback.py

Move functions out of xla.py closer to their users.

2023-08-08 14:40:42 -07:00

jet.py

Use lower-case PEP 585 names for types.

2023-06-23 15:12:14 -07:00

maps.py

Copybara import of the project:

2023-05-24 10:35:37 -07:00

mesh_utils.py

Allow device mesh handler to return None and use the default logic

2023-08-07 12:50:43 -07:00

multihost_utils.py

[JAX] Remove code that sets or tests --jax_coordination_service.

2023-07-28 13:13:13 -07:00

ode.py

Remove leading underscores in jax._src.numpy.util

2023-03-13 12:18:36 -07:00

pjit.py

Prune some exports from jax.experimental.pjit.

2023-08-01 13:27:17 -07:00

profiler.py

Add get_profiled_instructions_proto API in jax.experimental.profiler to get profiled instructions proto.

2023-06-29 21:02:05 -07:00

rnn.py

Use lower-case PEP 585 names for types.

2023-06-23 15:12:14 -07:00

serialize_executable.py

Stop passing CompileOptions when deserializing.

2023-05-10 16:22:54 -07:00

shard_map.py

[shmap-partial-auto] start adding partial auto / partial manual to shmap

2023-08-01 14:51:07 -07:00

topologies.py

Apply pyupgrade --py39-plus.

2023-07-21 14:49:44 -04:00

x64_context.py

Change JAX's copyright attribution from "Google LLC" to "The JAX Authors.".

2022-09-22 12:27:19 -07:00