Parker Schuh 261ff9e9ed Stop passing CompileOptions when deserializing. ...

PiperOrigin-RevId: 531034200

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

..

array_serialization

Instrument a new metric to measure the savings of async checkpoint in JAX.

2023-05-03 16:36:54 -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

gda_serialization

Fix typo in warnings in serialization functions.

2023-04-07 09:24:07 -07:00

jax2tf

[shape_poly] Remove old test limitations

2023-05-10 13:38:24 +02:00

sparse

Make dimensions argument of bcoo_reshape optional.

2023-05-09 10:38:18 -07:00

__init__.py

Move PartitionSpec into its own file (jax/_src/partition_spec.py).

2023-04-06 11:43:25 -07:00

checkify.py

Document why 'import name as name' is used

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

custom_partitioning.py

Remove tupling of custom call results.

2023-05-01 09:02:14 -07:00

host_callback.py

Change to simpler import for jax.config

2023-04-21 11:51:22 -07:00

jet.py

Split sharding_impls into its own Bazel target.

2023-04-10 10:15:58 -07:00

maps.py

Remove _PositionalSemantics class since it is not used anymore because jax.Array always has GLOBAL semantics

2023-03-17 13:30:04 -07:00

mesh_utils.py

Improve error message in _create_device_mesh_for_nd_torus

2023-04-28 03:15:06 -07:00

multihost_utils.py

Handle multihost pmap in pmap shard_map merge. This involves lifting the host local inputs to global inputs and vice-versa on the outputs.

2023-05-02 16:53:22 -07:00

ode.py

Remove leading underscores in jax._src.numpy.util

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

pjit.py

Move parse_flatten_op_sharding to sharding_impls.py to remove local import of pjit using that function from pxla.py

2023-04-11 19:26:25 -07:00

rnn.py

Make rnn_bwd_abstract_eval backwards compatible by guarding it agains the jaxlib version

2023-05-03 19:28:42 -07:00

serialize_executable.py

Stop passing CompileOptions when deserializing.

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

shard_map.py

[shard-map] experiment with lowering to a Call with attrs

2023-05-10 13:14:04 -07:00

topologies.py

Add an experimental API that allows compiling AOT for TPUs.

2023-04-19 13:33:59 -07:00

x64_context.py

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

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