Yash Katariya 6f770c936e Make call_tf's abstract eval rule effectful so that JAX's DCE pass doesn't prune the arguments passed to it. ...

PiperOrigin-RevId: 496022997

2022-12-16 20:42:54 -08:00

..

compilation_cache

Add c++ pickling for all the jax Sharding types.

2022-11-02 17:28:38 -07: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

Make call_tf's abstract eval rule effectful so that JAX's DCE pass doesn't prune the arguments passed to it.

2022-12-16 20:42:54 -08:00

sparse

[sparse] Use SparseInfo for both bcoo and bcsr format.

2022-12-16 15:00:38 -08:00

__init__.py

Document why 'import name as name' is used

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

callback.py

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

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

checkify.py

Document why 'import name as name' is used

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

custom_partitioning.py

(NFC) Prepare for migration from producing MHLO to producing StableHLO

2022-12-15 21:00:07 -08:00

global_device_array.py

Remove local imports of array.py. The remaining local imports are in pxla.py but I will chip away at them when we delete SDA and move some more APIs out of experimental.

2022-11-30 15:26:03 -08:00

host_callback.py

(NFC) Prepare for migration from producing MHLO to producing StableHLO

2022-12-15 21:00:07 -08:00

jet.py

BUG: fix jet rule for dynamic_slice

2022-11-02 21:37:41 -07:00

maps.py

[jax2tf] An alternative support for shape polymorphism for native serialization.

2022-12-08 08:19:35 +02:00

mesh_utils.py

Generalize TPU mesh computations.

2022-11-21 03:22:24 -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

Merge pull request #11077 from DanPuzzuoli:ode_dt_max

2022-10-28 16:05:25 -07: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