…module (#168221)"

This reverts commit cb3754f.

Reverting this change as it affects the import path of a publicly
used API.

ghstack-source-id: d68ea51
Pull Request resolved: #170615

…module (#168221)"

This reverts commit cb3754f.

Reverting this change as it affects the import path of a publicly
used API.

ghstack-source-id: 8c00738
Pull Request resolved: #170615

…module (#168221)"

This reverts commit cb3754f.

Reverting this change as it affects the import path of a publicly
used API.

ghstack-source-id: 29c372c
Pull Request resolved: #170615

…module (#168221)"

This reverts commit cb3754f.

Reverting this change as it affects the import path of a publicly
used API.

ghstack-source-id: 97ee16a
Pull Request resolved: #170615

### Motivation
We find it is too difficult to implement sharding strategies for DTensor, and this slows progress towards full operator coverage, and increases likelihood of bugs in sharding rules.  We also expect to add new Placement types (e.g. more complete support for StridedShard in the near term, but possibly others as well), and the current formulation of sharding strategies is not scalable to adding new placement types.

A primary reason it's so difficult to write sharding rules today is that they combine 2 things: (a) a mathematical description of which in/out shardings are correct for the op, (b) premature runtime optimization to avoid considering too many combinations over N-D mesh.

The proposal is to remove (b) and focus just on (a).

### tl;dr
1. Write sharding prop rules in terms of 1 mesh dim and using a 'placeholder' for generic sharding types
e.g. matmul rule returns:
```
[
   ShardPlaceholder(0), Replicate() -> ShardPlaceholder(0),
   Replicate(), ShardPlaceholder(1), -> ShardPlaceholder(1),
   ShardPlaceholder(1), ShardPlaceholder(0) -> Partial()
]
```
2. After registration, each rule gets automatically expanded to include the real sharding types discovered in the inputs at runtime, and add 'full replication' rule.
e.g.
if inputs are fully replicated, we drop the placeholder rules and only use
```
[
   Replicate(), Replicate() -> Replicate()
]
```
if 'Shard' is discovered in inputs, we fill placeholders like
```
[
   Shard(0), Replicate() -> Shard(0),
   Replicate(), Shard(1), -> Shard(1),
   Shard(1), Shard(0) -> Partial(),
   Replicate(), Replicate() -> Replicate()
]
```
if 'Shard' and 'StridedShard' are both discovered in the inputs, we expand to
```
[
   Shard(0), Replicate() -> Shard(0),
   Replicate(), Shard(1), -> Shard(1),
   Shard(1), Shard(0) -> Partial(),
   StridedShard(0), Replicate() -> StridedShard(0),
   Replicate(), StridedShard(1), -> StridedShard(1),
   StridedShard(1), StridedShard(0) -> Partial(),
   Replicate(), Replicate() -> Replicate()
]
```
3. After filling the placeholders, we expand to N-D mesh and find the minimum cost
(a) full enumeration via itertools.product is implemented and gives exact parity with rules like 'einsum' today
(b) optimized solution, starting from input placements and iterating in the order of increasing cost until reaching a min-cost solution _without_ having to fully enumerate - under development/prototyping

### This PR
* defines a 'single_dim strategy' function and a ShardingPlaceholder
* adds a util for expanding a single_dim strategy into a regular strategy
* supports StridedShard automatically via ShardingPlaceholder expansion
* writes rules for mm and cat and uses unit tests to validate their expansion

### Next Steps (PR stack)
* Support pointwise and foreach ops in the single_dim infra
* Hook up single-dim strategies to sharding_prop (op registration)
* Start to use single_dim rules to replace existing rules
* Improve the runtime of searching the fully expanded strategy
* Explore using decomps together with single-dim rules to support more operators

Pull Request resolved: #167677
Approved by: https://github.com/weifengpy
ghstack dependencies: #170615

Co-authored-by: Pian Pawakapan <pianpwk@meta.com>

…module (#168221)"

This reverts commit cb3754f.

Reverting this change as it affects the import path of a publicly
used API.

ghstack-source-id: b787988
Pull Request resolved: #170615

### Motivation
We find it is too difficult to implement sharding strategies for DTensor, and this slows progress towards full operator coverage, and increases likelihood of bugs in sharding rules.  We also expect to add new Placement types (e.g. more complete support for StridedShard in the near term, but possibly others as well), and the current formulation of sharding strategies is not scalable to adding new placement types.

A primary reason it's so difficult to write sharding rules today is that they combine 2 things: (a) a mathematical description of which in/out shardings are correct for the op, (b) premature runtime optimization to avoid considering too many combinations over N-D mesh.

The proposal is to remove (b) and focus just on (a).

### tl;dr
1. Write sharding prop rules in terms of 1 mesh dim and using a 'placeholder' for generic sharding types
e.g. matmul rule returns:
```
[
   ShardPlaceholder(0), Replicate() -> ShardPlaceholder(0),
   Replicate(), ShardPlaceholder(1), -> ShardPlaceholder(1),
   ShardPlaceholder(1), ShardPlaceholder(0) -> Partial()
]
```
2. After registration, each rule gets automatically expanded to include the real sharding types discovered in the inputs at runtime, and add 'full replication' rule.
e.g.
if inputs are fully replicated, we drop the placeholder rules and only use
```
[
   Replicate(), Replicate() -> Replicate()
]
```
if 'Shard' is discovered in inputs, we fill placeholders like
```
[
   Shard(0), Replicate() -> Shard(0),
   Replicate(), Shard(1), -> Shard(1),
   Shard(1), Shard(0) -> Partial(),
   Replicate(), Replicate() -> Replicate()
]
```
if 'Shard' and 'StridedShard' are both discovered in the inputs, we expand to
```
[
   Shard(0), Replicate() -> Shard(0),
   Replicate(), Shard(1), -> Shard(1),
   Shard(1), Shard(0) -> Partial(),
   StridedShard(0), Replicate() -> StridedShard(0),
   Replicate(), StridedShard(1), -> StridedShard(1),
   StridedShard(1), StridedShard(0) -> Partial(),
   Replicate(), Replicate() -> Replicate()
]
```
3. After filling the placeholders, we expand to N-D mesh and find the minimum cost
(a) full enumeration via itertools.product is implemented and gives exact parity with rules like 'einsum' today
(b) optimized solution, starting from input placements and iterating in the order of increasing cost until reaching a min-cost solution _without_ having to fully enumerate - under development/prototyping

### This PR
* defines a 'single_dim strategy' function and a ShardingPlaceholder
* adds a util for expanding a single_dim strategy into a regular strategy
* supports StridedShard automatically via ShardingPlaceholder expansion
* writes rules for mm and cat and uses unit tests to validate their expansion

### Next Steps (PR stack)
* Support pointwise and foreach ops in the single_dim infra
* Hook up single-dim strategies to sharding_prop (op registration)
* Start to use single_dim rules to replace existing rules
* Improve the runtime of searching the fully expanded strategy
* Explore using decomps together with single-dim rules to support more operators

Pull Request resolved: #167677
Approved by: https://github.com/weifengpy
ghstack dependencies: #170615

Co-authored-by: Pian Pawakapan <pianpwk@meta.com>

This PR adds the register_single_dim_strategy util,  and hooks it up to sharding_propagator.  It also tests the registration.

Notes:
* I didn't yet decide how multiple registrations should be handled.  I was planning to make it an error if you register twice for the same op for either single_dim or regular strategies.
* I took the cleanest path of integration for now in sharding_prop, reusing as much code as possible with the existing 'op_strategy' case.  I may have to change this later when integrating find_min_cost

Pull Request resolved: #170359
Approved by: https://github.com/weifengpy
ghstack dependencies: #170615, #167677

Co-authored-by: Pian Pawakapan <pianpwk@meta.com>

This reverts commit 32d0782.

Reverted pytorch#170359 on behalf of https://github.com/jeanschmidt due to Required to revert pytorch#167677 that is required to revert pytorch#170615 that is required to revert pytorch#170030 ([comment](pytorch#170359 (comment)))

This reverts commit c3e628e.

Reverted pytorch#167677 on behalf of https://github.com/jeanschmidt due to Required to revert pytorch#170615 that is required to rever pytorch#170030 ([comment](pytorch#167677 (comment)))

… module (pytorch#168221)" (pytorch#170615)

This reverts commit c65f67b.

Reverted pytorch#170615 on behalf of https://github.com/jeanschmidt due to Required to revert pytorch#170030 ([comment](pytorch#170615 (comment)))

…module (pytorch#168221)" (pytorch#170615)

This reverts commit cb3754f.

Reverting this change as it affects the import path of a publicly
used API.
Pull Request resolved: pytorch#170615
Approved by: https://github.com/wdvr, https://github.com/malfet

### Motivation
We find it is too difficult to implement sharding strategies for DTensor, and this slows progress towards full operator coverage, and increases likelihood of bugs in sharding rules.  We also expect to add new Placement types (e.g. more complete support for StridedShard in the near term, but possibly others as well), and the current formulation of sharding strategies is not scalable to adding new placement types.

A primary reason it's so difficult to write sharding rules today is that they combine 2 things: (a) a mathematical description of which in/out shardings are correct for the op, (b) premature runtime optimization to avoid considering too many combinations over N-D mesh.

The proposal is to remove (b) and focus just on (a).

### tl;dr
1. Write sharding prop rules in terms of 1 mesh dim and using a 'placeholder' for generic sharding types
e.g. matmul rule returns:
```
[
   ShardPlaceholder(0), Replicate() -> ShardPlaceholder(0),
   Replicate(), ShardPlaceholder(1), -> ShardPlaceholder(1),
   ShardPlaceholder(1), ShardPlaceholder(0) -> Partial()
]
```
2. After registration, each rule gets automatically expanded to include the real sharding types discovered in the inputs at runtime, and add 'full replication' rule.
e.g.
if inputs are fully replicated, we drop the placeholder rules and only use
```
[
   Replicate(), Replicate() -> Replicate()
]
```
if 'Shard' is discovered in inputs, we fill placeholders like
```
[
   Shard(0), Replicate() -> Shard(0),
   Replicate(), Shard(1), -> Shard(1),
   Shard(1), Shard(0) -> Partial(),
   Replicate(), Replicate() -> Replicate()
]
```
if 'Shard' and 'StridedShard' are both discovered in the inputs, we expand to
```
[
   Shard(0), Replicate() -> Shard(0),
   Replicate(), Shard(1), -> Shard(1),
   Shard(1), Shard(0) -> Partial(),
   StridedShard(0), Replicate() -> StridedShard(0),
   Replicate(), StridedShard(1), -> StridedShard(1),
   StridedShard(1), StridedShard(0) -> Partial(),
   Replicate(), Replicate() -> Replicate()
]
```
3. After filling the placeholders, we expand to N-D mesh and find the minimum cost
(a) full enumeration via itertools.product is implemented and gives exact parity with rules like 'einsum' today
(b) optimized solution, starting from input placements and iterating in the order of increasing cost until reaching a min-cost solution _without_ having to fully enumerate - under development/prototyping

### This PR
* defines a 'single_dim strategy' function and a ShardingPlaceholder
* adds a util for expanding a single_dim strategy into a regular strategy
* supports StridedShard automatically via ShardingPlaceholder expansion
* writes rules for mm and cat and uses unit tests to validate their expansion

### Next Steps (PR stack)
* Support pointwise and foreach ops in the single_dim infra
* Hook up single-dim strategies to sharding_prop (op registration)
* Start to use single_dim rules to replace existing rules
* Improve the runtime of searching the fully expanded strategy
* Explore using decomps together with single-dim rules to support more operators

Pull Request resolved: pytorch#167677
Approved by: https://github.com/weifengpy
ghstack dependencies: pytorch#170615

Co-authored-by: Pian Pawakapan <pianpwk@meta.com>

This PR adds the register_single_dim_strategy util,  and hooks it up to sharding_propagator.  It also tests the registration.

Notes:
* I didn't yet decide how multiple registrations should be handled.  I was planning to make it an error if you register twice for the same op for either single_dim or regular strategies.
* I took the cleanest path of integration for now in sharding_prop, reusing as much code as possible with the existing 'op_strategy' case.  I may have to change this later when integrating find_min_cost

Pull Request resolved: pytorch#170359
Approved by: https://github.com/weifengpy
ghstack dependencies: pytorch#170615, pytorch#167677

Co-authored-by: Pian Pawakapan <pianpwk@meta.com>

Enforce tensor_meta is not none for new single-dim rules.

Allow tensor_meta to continue to be None for existing rules for now. We
should consider in the future asserting tensor_meta is required in
DTensorSpec, but for now we just try to limit the bleeding.
Pull Request resolved: #170827
Approved by: https://github.com/dolpm
ghstack dependencies: #170615, #167677, #170359

This reverts commit 32d0782.

Reverted #170359 on behalf of https://github.com/jeanschmidt due to Required to revert #167677 that is required to revert #170615 that is required to revert #170030 ([comment](#170359 (comment)))

This reverts commit c3e628e.

Reverted #167677 on behalf of https://github.com/jeanschmidt due to Required to revert #170615 that is required to rever #170030 ([comment](#167677 (comment)))

… module (#168221)" (#170615)

This reverts commit c65f67b.

Reverted #170615 on behalf of https://github.com/jeanschmidt due to Required to revert #170030 ([comment](#170615 (comment)))

…module (#168221)" (#170615)

This reverts commit cb3754f.

Reverting this change as it affects the import path of a publicly
used API.
Pull Request resolved: #170615
Approved by: https://github.com/wdvr, https://github.com/malfet

### Motivation
We find it is too difficult to implement sharding strategies for DTensor, and this slows progress towards full operator coverage, and increases likelihood of bugs in sharding rules.  We also expect to add new Placement types (e.g. more complete support for StridedShard in the near term, but possibly others as well), and the current formulation of sharding strategies is not scalable to adding new placement types.

A primary reason it's so difficult to write sharding rules today is that they combine 2 things: (a) a mathematical description of which in/out shardings are correct for the op, (b) premature runtime optimization to avoid considering too many combinations over N-D mesh.

The proposal is to remove (b) and focus just on (a).

### tl;dr
1. Write sharding prop rules in terms of 1 mesh dim and using a 'placeholder' for generic sharding types
e.g. matmul rule returns:
```
[
   ShardPlaceholder(0), Replicate() -> ShardPlaceholder(0),
   Replicate(), ShardPlaceholder(1), -> ShardPlaceholder(1),
   ShardPlaceholder(1), ShardPlaceholder(0) -> Partial()
]
```
2. After registration, each rule gets automatically expanded to include the real sharding types discovered in the inputs at runtime, and add 'full replication' rule.
e.g.
if inputs are fully replicated, we drop the placeholder rules and only use
```
[
   Replicate(), Replicate() -> Replicate()
]
```
if 'Shard' is discovered in inputs, we fill placeholders like
```
[
   Shard(0), Replicate() -> Shard(0),
   Replicate(), Shard(1), -> Shard(1),
   Shard(1), Shard(0) -> Partial(),
   Replicate(), Replicate() -> Replicate()
]
```
if 'Shard' and 'StridedShard' are both discovered in the inputs, we expand to
```
[
   Shard(0), Replicate() -> Shard(0),
   Replicate(), Shard(1), -> Shard(1),
   Shard(1), Shard(0) -> Partial(),
   StridedShard(0), Replicate() -> StridedShard(0),
   Replicate(), StridedShard(1), -> StridedShard(1),
   StridedShard(1), StridedShard(0) -> Partial(),
   Replicate(), Replicate() -> Replicate()
]
```
3. After filling the placeholders, we expand to N-D mesh and find the minimum cost
(a) full enumeration via itertools.product is implemented and gives exact parity with rules like 'einsum' today
(b) optimized solution, starting from input placements and iterating in the order of increasing cost until reaching a min-cost solution _without_ having to fully enumerate - under development/prototyping

### This PR
* defines a 'single_dim strategy' function and a ShardingPlaceholder
* adds a util for expanding a single_dim strategy into a regular strategy
* supports StridedShard automatically via ShardingPlaceholder expansion
* writes rules for mm and cat and uses unit tests to validate their expansion

### Next Steps (PR stack)
* Support pointwise and foreach ops in the single_dim infra
* Hook up single-dim strategies to sharding_prop (op registration)
* Start to use single_dim rules to replace existing rules
* Improve the runtime of searching the fully expanded strategy
* Explore using decomps together with single-dim rules to support more operators

Pull Request resolved: #167677
Approved by: https://github.com/weifengpy
ghstack dependencies: #170615

Co-authored-by: Pian Pawakapan <pianpwk@meta.com>

This PR adds the register_single_dim_strategy util,  and hooks it up to sharding_propagator.  It also tests the registration.

Notes:
* I didn't yet decide how multiple registrations should be handled.  I was planning to make it an error if you register twice for the same op for either single_dim or regular strategies.
* I took the cleanest path of integration for now in sharding_prop, reusing as much code as possible with the existing 'op_strategy' case.  I may have to change this later when integrating find_min_cost

Pull Request resolved: #170359
Approved by: https://github.com/weifengpy
ghstack dependencies: #170615, #167677

Co-authored-by: Pian Pawakapan <pianpwk@meta.com>

Enforce tensor_meta is not none for new single-dim rules.

Allow tensor_meta to continue to be None for existing rules for now. We
should consider in the future asserting tensor_meta is required in
DTensorSpec, but for now we just try to limit the bleeding.
Pull Request resolved: #170827
Approved by: https://github.com/dolpm
ghstack dependencies: #170615, #167677, #170359

…module (pytorch#168221)" (pytorch#170615)

This reverts commit cb3754f.

Reverting this change as it affects the import path of a publicly
used API.
Pull Request resolved: pytorch#170615
Approved by: https://github.com/wdvr, https://github.com/malfet

### Motivation
We find it is too difficult to implement sharding strategies for DTensor, and this slows progress towards full operator coverage, and increases likelihood of bugs in sharding rules.  We also expect to add new Placement types (e.g. more complete support for StridedShard in the near term, but possibly others as well), and the current formulation of sharding strategies is not scalable to adding new placement types.

A primary reason it's so difficult to write sharding rules today is that they combine 2 things: (a) a mathematical description of which in/out shardings are correct for the op, (b) premature runtime optimization to avoid considering too many combinations over N-D mesh.

The proposal is to remove (b) and focus just on (a).

### tl;dr
1. Write sharding prop rules in terms of 1 mesh dim and using a 'placeholder' for generic sharding types
e.g. matmul rule returns:
```
[
   ShardPlaceholder(0), Replicate() -> ShardPlaceholder(0),
   Replicate(), ShardPlaceholder(1), -> ShardPlaceholder(1),
   ShardPlaceholder(1), ShardPlaceholder(0) -> Partial()
]
```
2. After registration, each rule gets automatically expanded to include the real sharding types discovered in the inputs at runtime, and add 'full replication' rule.
e.g.
if inputs are fully replicated, we drop the placeholder rules and only use
```
[
   Replicate(), Replicate() -> Replicate()
]
```
if 'Shard' is discovered in inputs, we fill placeholders like
```
[
   Shard(0), Replicate() -> Shard(0),
   Replicate(), Shard(1), -> Shard(1),
   Shard(1), Shard(0) -> Partial(),
   Replicate(), Replicate() -> Replicate()
]
```
if 'Shard' and 'StridedShard' are both discovered in the inputs, we expand to
```
[
   Shard(0), Replicate() -> Shard(0),
   Replicate(), Shard(1), -> Shard(1),
   Shard(1), Shard(0) -> Partial(),
   StridedShard(0), Replicate() -> StridedShard(0),
   Replicate(), StridedShard(1), -> StridedShard(1),
   StridedShard(1), StridedShard(0) -> Partial(),
   Replicate(), Replicate() -> Replicate()
]
```
3. After filling the placeholders, we expand to N-D mesh and find the minimum cost
(a) full enumeration via itertools.product is implemented and gives exact parity with rules like 'einsum' today
(b) optimized solution, starting from input placements and iterating in the order of increasing cost until reaching a min-cost solution _without_ having to fully enumerate - under development/prototyping

### This PR
* defines a 'single_dim strategy' function and a ShardingPlaceholder
* adds a util for expanding a single_dim strategy into a regular strategy
* supports StridedShard automatically via ShardingPlaceholder expansion
* writes rules for mm and cat and uses unit tests to validate their expansion

### Next Steps (PR stack)
* Support pointwise and foreach ops in the single_dim infra
* Hook up single-dim strategies to sharding_prop (op registration)
* Start to use single_dim rules to replace existing rules
* Improve the runtime of searching the fully expanded strategy
* Explore using decomps together with single-dim rules to support more operators

Pull Request resolved: pytorch#167677
Approved by: https://github.com/weifengpy
ghstack dependencies: pytorch#170615

Co-authored-by: Pian Pawakapan <pianpwk@meta.com>

This reverts commit 32d0782.

Reverted pytorch#170359 on behalf of https://github.com/jeanschmidt due to Required to revert pytorch#167677 that is required to revert pytorch#170615 that is required to revert pytorch#170030 ([comment](pytorch#170359 (comment)))

This reverts commit c3e628e.

Reverted pytorch#167677 on behalf of https://github.com/jeanschmidt due to Required to revert pytorch#170615 that is required to rever pytorch#170030 ([comment](pytorch#167677 (comment)))

… module (pytorch#168221)" (pytorch#170615)

This reverts commit c65f67b.

Reverted pytorch#170615 on behalf of https://github.com/jeanschmidt due to Required to revert pytorch#170030 ([comment](pytorch#170615 (comment)))

…module (pytorch#168221)" (pytorch#170615)

This reverts commit cb3754f.

Reverting this change as it affects the import path of a publicly
used API.
Pull Request resolved: pytorch#170615
Approved by: https://github.com/wdvr, https://github.com/malfet

### Motivation
We find it is too difficult to implement sharding strategies for DTensor, and this slows progress towards full operator coverage, and increases likelihood of bugs in sharding rules.  We also expect to add new Placement types (e.g. more complete support for StridedShard in the near term, but possibly others as well), and the current formulation of sharding strategies is not scalable to adding new placement types.

A primary reason it's so difficult to write sharding rules today is that they combine 2 things: (a) a mathematical description of which in/out shardings are correct for the op, (b) premature runtime optimization to avoid considering too many combinations over N-D mesh.

The proposal is to remove (b) and focus just on (a).

### tl;dr
1. Write sharding prop rules in terms of 1 mesh dim and using a 'placeholder' for generic sharding types
e.g. matmul rule returns:
```
[
   ShardPlaceholder(0), Replicate() -> ShardPlaceholder(0),
   Replicate(), ShardPlaceholder(1), -> ShardPlaceholder(1),
   ShardPlaceholder(1), ShardPlaceholder(0) -> Partial()
]
```
2. After registration, each rule gets automatically expanded to include the real sharding types discovered in the inputs at runtime, and add 'full replication' rule.
e.g.
if inputs are fully replicated, we drop the placeholder rules and only use
```
[
   Replicate(), Replicate() -> Replicate()
]
```
if 'Shard' is discovered in inputs, we fill placeholders like
```
[
   Shard(0), Replicate() -> Shard(0),
   Replicate(), Shard(1), -> Shard(1),
   Shard(1), Shard(0) -> Partial(),
   Replicate(), Replicate() -> Replicate()
]
```
if 'Shard' and 'StridedShard' are both discovered in the inputs, we expand to
```
[
   Shard(0), Replicate() -> Shard(0),
   Replicate(), Shard(1), -> Shard(1),
   Shard(1), Shard(0) -> Partial(),
   StridedShard(0), Replicate() -> StridedShard(0),
   Replicate(), StridedShard(1), -> StridedShard(1),
   StridedShard(1), StridedShard(0) -> Partial(),
   Replicate(), Replicate() -> Replicate()
]
```
3. After filling the placeholders, we expand to N-D mesh and find the minimum cost
(a) full enumeration via itertools.product is implemented and gives exact parity with rules like 'einsum' today
(b) optimized solution, starting from input placements and iterating in the order of increasing cost until reaching a min-cost solution _without_ having to fully enumerate - under development/prototyping

### This PR
* defines a 'single_dim strategy' function and a ShardingPlaceholder
* adds a util for expanding a single_dim strategy into a regular strategy
* supports StridedShard automatically via ShardingPlaceholder expansion
* writes rules for mm and cat and uses unit tests to validate their expansion

### Next Steps (PR stack)
* Support pointwise and foreach ops in the single_dim infra
* Hook up single-dim strategies to sharding_prop (op registration)
* Start to use single_dim rules to replace existing rules
* Improve the runtime of searching the fully expanded strategy
* Explore using decomps together with single-dim rules to support more operators

Pull Request resolved: pytorch#167677
Approved by: https://github.com/weifengpy
ghstack dependencies: pytorch#170615

Co-authored-by: Pian Pawakapan <pianpwk@meta.com>

This PR adds the register_single_dim_strategy util,  and hooks it up to sharding_propagator.  It also tests the registration.

Notes:
* I didn't yet decide how multiple registrations should be handled.  I was planning to make it an error if you register twice for the same op for either single_dim or regular strategies.
* I took the cleanest path of integration for now in sharding_prop, reusing as much code as possible with the existing 'op_strategy' case.  I may have to change this later when integrating find_min_cost

Pull Request resolved: pytorch#170359
Approved by: https://github.com/weifengpy
ghstack dependencies: pytorch#170615, pytorch#167677

Co-authored-by: Pian Pawakapan <pianpwk@meta.com>

Enforce tensor_meta is not none for new single-dim rules.

Allow tensor_meta to continue to be None for existing rules for now. We
should consider in the future asserting tensor_meta is required in
DTensorSpec, but for now we just try to limit the bleeding.
Pull Request resolved: pytorch#170827
Approved by: https://github.com/dolpm
ghstack dependencies: pytorch#170615, pytorch#167677, pytorch#170359

…module (#168221)"

This reverts commit cb3754f.

Reverting this change as it affects the import path of a publicly
used API.

ghstack-source-id: b787988
Pull Request resolved: pytorch/pytorch#170615