Yash Katariya
14451492c9
Delete OpShardingSharding export since it has been 3 months since it was deprecated. Also remove deprecation warnings for MeshPspecSharding.
...
PiperOrigin-RevId: 538880293
2023-06-08 13:53:38 -07:00
Peter Hawkins
b4402185db
Move PartitionSpec into its own file (jax/_src/partition_spec.py).
...
No functional changes intended.
A subsequent change will move ParsedPartitionSpec and array mapping utilities here also.
PiperOrigin-RevId: 522393166
2023-04-06 11:43:25 -07:00
Yash Katariya
7442faa715
Remove MeshPspecSharding since it has been more than 3 months since it was deprecated (Nov 2, 2022). The replacement name is NamedSharding.
...
PiperOrigin-RevId: 520072687
2023-03-28 10:47:42 -07:00
Peter Hawkins
1925aa1109
Split Sharding subclasses out of _src/sharding.py into _src/sharding_impls.py
...
By defining the Sharding base class in its own module, we can pull it out into a separate Bazel submodule, which will help pytype inference when defining Array.
PiperOrigin-RevId: 516223009
2023-03-13 08:50:18 -07:00
Yash Katariya
0ffdeb3de2
Rename jax.sharding.OpShardingSharding to jax.sharding.GSPMDSharding. jax.sharding.OpShardingSharding will be removed in 3 months from Feb 17, 2023.
...
PiperOrigin-RevId: 510556189
2023-02-17 17:11:06 -08:00
Peter Hawkins
8268cd562d
Add infrastructure for managing deprecations.
...
Use it to deprecate jax.experimental.PartitionSpec, jax.interpreters.pxla.PartitionSpec, jax.interpreters.pxla.Mesh.
PiperOrigin-RevId: 508349776
2023-02-09 05:48:40 -08:00
Jake VanderPlas
26f2f97805
Document why 'import name as name' is used
2022-12-14 15:07:04 -08:00
Yash Katariya
934bc4e1b3
Move PartitionSpec and Mesh out of experimental and into the sharding namespace. The new API endpoint is jax.sharding.PartitionSpec and jax.sharding.Mesh.
...
PiperOrigin-RevId: 492358238
2022-12-01 19:28:32 -08:00
Yash Katariya
cc5af7ed98
Rename ReshapeableDevicesSharding to PositionalSharding and add an alias NamedSharding for MeshPspecSharding.
...
`MeshPspecSharding` name will be replaced with `NamedSharding` in 3 months.
PiperOrigin-RevId: 485753078
2022-11-02 19:13:13 -07:00
Matthew Johnson
95eb4249bb
tweaks to DevicesSharding
...
1. rename DevicesSharding -> ReshapeableDevicesSharding
2. fix repr to print device order faithfully
3. respect shape of np.ndarray argument to __init__
2022-10-25 14:28:48 -07:00
Matthew Johnson
43098f906a
initial commit of DevicesSharding (fka SimpleSharding)
...
need to add tests!
Co-authored-by: Yash Katariya <yashkatariya@google.com>
Co-authored-by: Sharad Vikram <sharad.vikram@gmail.com>
2022-10-18 21:10:24 -07:00
Yash Katariya
9e4114f0f1
Move array.py and sharding.py from experimental/ to _src/.
...
PiperOrigin-RevId: 477201711
2022-09-27 10:06:52 -07:00
