Yash Katariya 7fbf8ec669 Fix Forward. The fix is on the user's end. Original PR: https://github.com/google/jax/pull/12217 ...

Co-authored-by: Matthew Johnson <mattjj@google.com>
Co-authored-by: Yash Katariya <yashkatariya@google.com>
PiperOrigin-RevId: 472999907

2022-09-08 08:49:40 -07:00

..

compilation_cache

Reapply: Use MLIR bytecode when passing IR to backends.

2022-09-08 08:11:16 -07:00

gda_serialization

Fix some cases where deserialization hangs:

2022-09-06 10:27:57 -07:00

jax2tf

Merge pull request #12251 from gnecula:tf_refactor2

2022-09-07 06:15:22 -07:00

sparse

[sparse] add sparse support for dynamic_slice

2022-09-01 13:42:02 -07:00

__init__.py

Delete sharded_jit

2022-07-29 08:19:52 -07:00

array.py

Fix Forward. The fix is on the user's end. Original PR: https://github.com/google/jax/pull/12217

2022-09-08 08:49:40 -07:00

callback.py

rolling forward deletion of custom_jvp_call_jaxpr_p yet again...

2022-08-18 14:02:40 -07:00

checkify.py

CI: fix flake8 ignore declarations

2022-04-21 13:44:12 -07:00

global_device_array.py

Add support to handle arbitrary shardings to KeyArray. Resolve all the TODOs that were created before.

2022-08-31 22:54:06 -07:00

host_callback.py

rolling forward deletion of custom_jvp_call_jaxpr_p yet again...

2022-08-18 14:02:40 -07:00

jet.py

Roll forward: Upgrade logistic into a primitive.

2022-09-07 06:06:56 -07:00

maps.py

Enable jax_array on CPU for the entire JAX test suite!

2022-08-19 10:04:35 -07:00

mesh_utils.py

[mesh_utils] Avoid relying on process-tiled device order

2022-06-06 13:30:56 -07:00

multihost_utils.py

[JAX] Deprecate .to_py() property on arrays. Implement __array__ instead.

2022-08-25 07:28:27 -07:00

ode.py

[x64] make ode_test pass with strict dtype promotion

2022-06-17 15:21:10 -07:00

pjit.py

Make jit == pjit. This means that the lowering and execution paths of jit and pjit are merged.

2022-08-29 22:03:21 -07:00

sharding.py

Add __hash__ and __eq__ to PmapSharding

2022-08-31 14:27:16 -07:00

x64_context.py

[x64] make x64_context tests work with updated default dtypes

2021-12-08 13:05:40 -08:00