Yash Katariya 94f0ccc54a Fix host_callback for pjit which was using REPLICATED which was a CanonicalizedParsedPspec ...

PiperOrigin-RevId: 501713533

2023-01-12 18:00:33 -08:00

..

compilation_cache

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

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

gda_serialization

Make MeshPspecSharding an alias for NamedSharding (it was the other way around before this CL).

2022-11-14 14:44:00 -08:00

jax2tf

Allow call_tf to accept shape-polymorphic inputs if the output shapes are fully static.

2023-01-12 13:04:39 -08:00

sparse

[sparse] generalize batch rule for bcoo_sum_duplicates & improve tests

2023-01-12 14:28:12 -08:00

__init__.py

Move functions into api_util.py and dispatch.py to remove circular import error when pjit is imported in api.py for merging the jit and pjit frontend API.

2022-12-22 08:42:05 -08:00

callback.py

Move jax.linear_util to jax._src.linear_util

2022-12-20 14:49:27 -08:00

checkify.py

Document why 'import name as name' is used

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

custom_partitioning.py

cleanup: remove several unused imports across the package

2022-12-28 12:49:17 -08:00

global_device_array.py

Simplify Array's shard_arg_handler by merging pmap and pjit/xmap paths

2022-12-27 10:16:44 -08:00

host_callback.py

Fix host_callback for pjit which was using REPLICATED which was a CanonicalizedParsedPspec

2023-01-12 18:00:33 -08:00

jet.py

cleanup: remove several unused imports across the package

2022-12-28 12:49:17 -08:00

maps.py

[docs] donate_argnums FAQ link to rst format

2023-01-10 18:11:08 +09:00

mesh_utils.py

Update the non-contiguous error message to not say GDA anymore

2023-01-11 15:35:15 -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

Move jax.linear_util to jax._src.linear_util

2022-12-20 14:49:27 -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

Add CUDNN custom call for LSTM. Exposed as jax.experimental.rnn module.

2022-11-28 14:31:48 -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