Repository: pytorch
Issue: [DTensor] addmm strategy cannot handle 0-element input tensors
Labels: triaged,module: dtensor,release notes: distributed (dtensor)

### 🐛 Describe the bug

Here is a test example from `test_dtensor_ops.py` for `addmm` operator:
```
@with_comms
def test_addmm_broadcast(self):
        device_mesh = DeviceMesh(self.device_type, list(range(self.world_size)))
        shard_spec = [Shard(0)]
        replica_spec = [Replicate()]

        tensor_to_shard = torch.randn(12, 0)
        mat1 = distribute_tensor(tensor_to_shard, device_mesh, shard_spec)
        tensor_to_replicate = torch.randn(0, 4)
        mat2 = distribute_tensor(tensor_to_replicate, device_mesh, replica_spec)
        input_tensor = torch.randn(4)
        input = distribute_tensor(input_tensor, device_mesh, replica_spec)

        dist_res = torch.addmm(input, mat1, mat2)
        local_res = torch.addmm(input_tensor, tensor_to_shard, tensor_to_replicate)
        print(f"rank={self.rank}, dist_res._spec={dist_res._spec}, dist_res._local_tensor={dist_res._local_tensor}")
        print(f"rank={self.rank}, dist_res={dist_res.full_tensor()}, local_res={local_res}")
        self.assertNotEqual(dist_res.full_tensor(), local_res)  # this should be equal!!
```

Expected Behavior in PyTorch:
```
>>> x = torch.rand(4)
>>> x
tensor([0.6197, 0.5531, 0.4860, 0.7942])
>>> y = torch.rand(12, 0)
>>> z = torch.rand(0, 4)
>>> torch.addmm(x, y, z)
tensor([[0.6197, 0.5531, 0.4860, 0.7942],
        [0.6197, 0.5531, 0.4860, 0.7942],
        [0.6197, 0.5531, 0.4860, 0.7942],
        [0.6197, 0.5531, 0.4860, 0.7942],
        [0.6197, 0.5531, 0.4860, 0.7942],
        [0.6197, 0.5531, 0.4860, 0.7942],
        [0.6197, 0.5531, 0.4860, 0.7942],
        [0.6197, 0.5531, 0.4860, 0.7942],
        [0.6197, 0.5531, 0.4860, 0.7942],
        [0.6197, 0.5531, 0.4860, 0.7942],
        [0.6197, 0.5531, 0.4860, 0.7942],
        [0.6197, 0.5531, 0.4860, 0.7942]])
```

Actual Result:
```
[rank1]:[2024-01-17 19:55:59,111] torch.testing._internal.common_distributed: [ERROR]   File "/data/users/xilunwu/oss/pytorch/test/distributed/_tensor/test_matrix_ops.py", line 55, in test_addmm_broadcast
[rank1]:[2024-01-17 19:55:59,111] torch.testing._internal.common_distributed: [ERROR]     dist_res = torch.addmm(input, mat1, mat2)
[rank1]:[2024-01-17 19:55:59,111] torch.testing._internal.common_distributed: [ERROR]                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank1]:[2024-01-17 19:55:59,111] torch.testing._internal.common_distributed: [ERROR]   File "/data/users/xilunwu/oss/pytorch/torch/distributed/_tensor/api.py", line 289, in __torch_dispatch__
[rank1]:[2024-01-17 19:55:59,111] torch.testing._internal.common_distributed: [ERROR]     return DTensor._op_dispatcher.dispatch(
[rank1]:[2024-01-17 19:55:59,111] torch.testing._internal.common_distributed: [ERROR]            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank1]:[2024-01-17 19:55:59,111] torch.testing._internal.common_distributed: [ERROR]   File "/data/users/xilunwu/oss/pytorch/torch/distributed/_tensor/dispatch.py", line 113, in dispatch
[rank1]:[2024-01-17 19:55:59,111] torch.testing._internal.common_distributed: [ERROR]     self.sharding_propagator.propagate(op_info)
[rank1]:[2024-01-17 19:55:59,111] torch.testing._internal.common_distributed: [ERROR]   File "/data/users/xilunwu/oss/pytorch/torch/distributed/_tensor/sharding_prop.py", line 164, in propagate
[rank1]:[2024-01-17 19:55:59,111] torch.testing._internal.common_distributed: [ERROR]     output_sharding = self.propagate_op_sharding(op_info.schema)
[rank1]:[2024-01-17 19:55:59,111] torch.testing._internal.common_distributed: [ERROR]                       ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank1]:[2024-01-17 19:55:59,111] torch.testing._internal.common_distributed: [ERROR]   File "/data/users/xilunwu/oss/pytorch/torch/distributed/_tensor/sharding_prop.py", line 216, in propagate_op_sharding_non_cached
[rank1]:[2024-01-17 19:55:59,111] torch.testing._internal.common_distributed: [ERROR]     output_strategy = self._select_strategy(op_strategy)
[rank1]:[2024-01-17 19:55:59,111] torch.testing._internal.common_distributed: [ERROR]                       ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank1]:[2024-01-17 19:55:59,111] torch.testing._internal.common_distributed: [ERROR]   File "/data/users/xilunwu/oss/pytorch/torch/distributed/_tensor/sharding_prop.py", line 394, in _select_strategy
[rank1]:[2024-01-17 19:55:59,111] torch.testing._internal.common_distributed: [ERROR]     return strategy.strategies[strategy_costs.index(min(strategy_costs))]
[rank1]:[2024-01-17 19:55:59,111] torch.testing._internal.common_distributed: [ERROR]                                                     ^^^^^^^^^^^^^^^^^^^
[rank1]:[2024-01-17 19:55:59,111] torch.testing._internal.common_distributed: [ERROR] ValueError: min() arg is an empty sequence
```
This log means there's no doable sharding for `addmm` on input `DTensor`s, which I believe should not be true.

Solution:
change the util function `is_tensor_shardable` to allow sharding on empty `DTensor`s. Check out PR: #117726 

New Error:
1. the result `DTensor` spec seems problematic. After `addmm`, the result `DTensor` has `_local_tensor=[ 0.0461,  0.4024, -1.0115,  0.2167]` and `_spec=Spec(S(0) on (12, 4))`. This should be the result of broadcasting and looks almost correct. An ideal result would be `_local_tensor=[[ 0.0461,  0.4024, -1.0115,  0.2167], [ 0.0461,  0.4024, -1.0115,  0.2167], [ 0.0461,  0.4024, -1.0115,  0.2167]]` so that calling `full_tensor()` will give a [12, 4] tensor object.

2. The result of `full_tensor()` over `DTensor` whose `_local_tensor=[ 0.0461,  0.4024, -1.0115,  0.2167]` and `_spec=Spec(S(0) on (12, 4))` will be a [16] local tensor:
`[ 0.0461,  0.4024, -1.0115,  0.2167,  0.0461,  0.4024, -1.0115,  0.2167, 0.0461,  0.4024, -1.0115,  0.2167,  0.0461,  0.4024, -1.0115,  0.2167]`

### Versions

pytorch build from source with PR #117726

cc @wanchaol @tianyu-l