Peter Hawkins f68f1c0cd0 Upgrade logistic (sigmoid) function into a lax primitive. ...

This allows us to lower it to `mhlo.logistic`, which allows XLA to generate more efficient code.

PiperOrigin-RevId: 470300985

2022-08-26 11:58:28 -07:00

..

compilation_cache

[JAX] Prepare not to export jax._src by default.

2022-08-23 09:36:47 -07:00

gda_serialization

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

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

jax2tf

Upgrade logistic (sigmoid) function into a lax primitive.

2022-08-26 11:58:28 -07:00

sparse

[sparse] add support for bcoo equivalent of lax.slice

2022-08-25 09:19:44 -07:00

__init__.py

Delete sharded_jit

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

array.py

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

2022-08-25 07:28:27 -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

Introduce class PyShardedBuffer that contains a vector of PjRtBuffer and cache it in GDA.

2022-08-25 10:32:39 -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

Upgrade logistic (sigmoid) function into a lax primitive.

2022-08-26 11:58:28 -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

Replace pspec with the original sharding that's already present from which pspec was extracted.

2022-08-17 19:46:39 -07:00

sharding.py

Add checkpointing support for Array similar to GDA.

2022-08-22 13:35:18 -07:00

x64_context.py

[x64] make x64_context tests work with updated default dtypes

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