Yuanzhong Xu 2e4687a62e PartitionSpec: allow partially specified sharding ...

PiperOrigin-RevId: 421603194

2022-01-13 10:35:15 -08:00

..

__init__.py

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

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

ad.py

fix a custom_vjp post_process bug, related cleanups

2022-01-12 07:51:50 -08:00

batching.py

fix a custom_vjp post_process bug, related cleanups

2022-01-12 07:51:50 -08:00

invertible_ad.py

Use imports relative to the jax package consistently, rather than .-relative imports.

2021-11-24 07:48:29 -08:00

masking.py

Move dtype canonicalization out of core.AbstractValue subclasses.

2021-12-07 06:13:07 -08:00

mlir.py

PartitionSpec: allow partially specified sharding

2022-01-13 10:35:15 -08:00

partial_eval.py

fix a custom_vjp post_process bug, related cleanups

2022-01-12 07:51:50 -08:00

pxla.py

Fix the index for fully replicated values. Each dimension of the aval should have an index. For example, if an aval has a shape (4, 4) then each dimension of the aval maybe sharded or replicated. Hence each dimension will have an index associated with it. For fully replicated avals, the dimensions are always replicated hence index is always slice(None).

2022-01-08 17:06:48 -08:00

sharded_jit.py

PartitionSpec: allow partially specified sharding

2022-01-13 10:35:15 -08:00

xla.py

PartitionSpec: allow partially specified sharding

2022-01-13 10:35:15 -08:00