Peter Hawkins 40c80d7d0a Remove jax._src from JAX namespace. ...

This is a JAX-internal name and not subject to any deprecation policy. Please avoid the use of JAX-internal functions outside JAX.

PiperOrigin-RevId: 473243243

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

..

compilation_cache

Reapply: Use MLIR bytecode when passing IR to backends.

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

gda_serialization

Support deserializing from non-chunked storage such as tensorstore.array.

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

jax2tf

Remove jax._src from JAX namespace.

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

sparse

Remove jax._src from JAX namespace.

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

__init__.py

Delete sharded_jit

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

array.py

Enable testArrayCopy now that its fixed.

2022-09-08 14:40:06 -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