Yash Katariya 83d7e3f085 Remove the fallback to lower_xla_callable that exist for jit(pmap) cases when Array was enabled and add minimal support to lower_sharding_computation. ...

The `jit(pmap)` codepath is added to `lower_sharding_computation` to delete the `lower_xla_callable` codepath when `jax.Array` is enabled by default. This will help in cleaning up the codebase and get rid of tech debt.

* Round trip through host for `Array`'s that have PmapSharding and come through the `jit` path (exactly like SDAs).

* For other cases i.e. when `num_replicas > 1`, default to the `_execute_replicated` path in dispatch.py from `lower_sharding_computation`. This is exactly same to what happens in `lower_xla_callable`.

PiperOrigin-RevId: 471033420

2022-08-30 10:46:23 -07:00

..

__init__.py

[JAX] Move traceback_util.py into jax._src.

2020-11-04 09:02:59 -08:00

ad.py

Merge pull request #11900 from sharadmv:for-loop

2022-08-18 20:14:17 -07:00

batching.py

Add spmd_axis_name to vmap to allow constraining mapped PartitionSpecs.

2022-08-08 19:41:42 -07:00

mlir.py

Migrate from MLIR-HLO's CHLO to StableHLO's CHLO

2022-08-26 09:35:23 -07:00

partial_eval.py

rolling forward deletion of custom_jvp_call_jaxpr_p yet again...

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

pxla.py

Remove the fallback to lower_xla_callable that exist for jit(pmap) cases when Array was enabled and add minimal support to lower_sharding_computation.

2022-08-30 10:46:23 -07:00

xla.py

BUG: ensure that boolean scalars are never marked weak

2022-07-07 15:41:23 -07:00