George Necula 7eb7baa6f1 [jax2tf] Add xmap tests ...

PiperOrigin-RevId: 505608747

2023-01-30 00:24:30 -08:00

..

compilation_cache

Add {SerializeAs,ParseFrom}String to xe.CompileOptions.

2023-01-03 09:15:34 -08:00

gda_serialization

Delete GDA serialization tests

2023-01-27 19:07:25 -08:00

jax2tf

[jax2tf] Add xmap tests

2023-01-30 00:24:30 -08:00

sparse

[sparse] Add function that fixes out-of-bound indices.

2023-01-24 11:46:46 -08:00

__init__.py

Add IO callback

2023-01-17 13:55:05 -08:00

checkify.py

Document why 'import name as name' is used

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

custom_partitioning.py

Fix trailing whiteshape + failing doc test + removing first section title

2023-01-24 15:55:17 -08:00

global_device_array.py

Delete the pjit 101 tutorial which is subsumed by https://jax.readthedocs.io/en/latest/notebooks/Distributed_arrays_and_automatic_parallelization.html

2023-01-20 20:35:47 -08:00

host_callback.py

Raise an error when attempting to mutate Jaxpr objects

2023-01-23 09:37:58 -08:00

jet.py

Add pjit support in jax.experimental.jet

2023-01-23 15:51:47 -08:00

maps.py

Make xmap use dispatch.sharded_lowering as dispatch.lower_xla_callable is deprecated.

2023-01-16 09:42:40 -08:00

mesh_utils.py

Delete the pjit 101 tutorial which is subsumed by https://jax.readthedocs.io/en/latest/notebooks/Distributed_arrays_and_automatic_parallelization.html

2023-01-20 20:35:47 -08:00

multihost_utils.py

Add typehints and point to the correct endpoint of Mesh and PartitionSpec in the args section.

2022-12-05 09:49:18 -08:00

ode.py

fix transcription error for initial_step_size (thanks @AlbertMitjans)

2023-01-14 21:10:44 -08:00

pjit.py

Move pjit.py to jax/_src in preparation for merging the jit and pjit frontend APIs

2022-12-16 13:07:15 -08:00

rnn.py

CI: fix mypy jaxlib version

2023-01-24 06:57:23 -08:00

serialize_executable.py

Refactor create_cpp_call to be a method on MeshExecutable

2022-11-17 18:05:10 -08:00

x64_context.py

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

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