Repository: pytorch
Issue: Support for variable number of dimensions in functorch.dim dims() and autonaming induced inconsistencies
Labels: triaged,module: functorch,module: first class dims

### 🚀 The feature, motivation and pitch

I propose to expand the api of functorch.dim.dims() with the possibility to pass a list of dimension names as an argument. 

### Describe current behavior

As of now, the syntax (dim_1, dim_2, ... dim_n) = dims(n) requires the user to pre-specify in advance the number of dimensions and their names. A call of the type dimensions = dims(n) results in a tuple dimensions = (d1, d2, ... dn) with the names d1, ... dn fixed internally withput any user choice. The code snippet illustrates the issue. cc torchdim contributors: @zdevito @ae99 

```
import torch
from functorch.dim import dims

# current behavior leads to confusing output.
batch_dims = dims(sizes = [10,5])
event_dims = dims(sizes = [3,2])

print(batch_dims)                           # produces (d0, d1)
print(event_dims)                           # produces (d0, d1)

print([dim.size for dim in batch_dims])     # produces [10,5]
print([dim.size for dim in event_dims])     # produces [3,2]

# This is confusing as there exist now two dimensions d0 with different properties

# The issue becomes even clearer, when using the dimensions to define tensors
# with first class dims:

some_dims = dims(2)
unrelated_dims = dims(2)

A = torch.randn([2,2])
A_fc_1 = A[some_dims]   
A_fc_2 = A[unrelated_dims]   

print(A_fc_1)       # produces tensor([[ ...]]) with dims=(d0,d1), sizes=(2,2)
print(A_fc_2)       # produces tensor([[ ...]]) with dims=(d0,d1), sizes=(2,2)

# this makes it look like there the same dimensions are indexing A_fc_1 and A_fc_2
# even though this is not the case and i am unsure about the behavior of e.g. the
# einsum-style ops.
```
It seems as if there is currently no possibility to use the dims() function to create a variable number of dimensions with userspecified names.

### Describe desired behavior

```
# The following is envisioned after resolving the issue
batch_dims = dims(sizes = [10,5], names = ['bd_1', 'bd_2'])
event_dims = dims(sizes = [3,2], names = ['ed_1', 'ed_2'])

print(batch_dims)    # produces (bd_1, bd_2)
print(event_dims)    # produces (ed_1, ed_2)

# This would allow for a construction that produces a variable number of batch_dims
# depending on user input.
bd_sizes = [2,3,4]
bd_names = ['bd_1', 'bd_2', 'bd_3']
batch_dims = dims(sizes = bd_sizes, names = bd_names)

# Where in the above, it is to be understood that bd_sizes and bd_names are reflecting
# the input of some user that specifies the dimensionality of a problem (e.g. if
# some offset should be applied to some 1d, 2d, 3d, tensor).
```

### Describe benefit of proposed changes

The current behavior is undesirable when the number of dimensions is not known beforehand as there is no method for naming individual dimensions when instantiating a variable number of dimensions at once. With the proposed api changes, it would be possible to instantiate and distinguish a variable number of dimensions. This would make it possible to create multiple dimensions at once without hardcoding their name and multiplicity.

### Alternatives

Named Tensors does not seem to be maintained anymore. Alternatives to creating variable-length tuples of dimensions with named entries do not seem possible with the current API for the functorch.dim.dims() function (apart from iteratively executing some string of code).

### Additional context

I work on probabilistic programming and would like to create uniquely identifiable or at least named batch/event dimensions by passing lists of names and dimsizes to dims().

cc @zou3519 @Chillee @samdow @kshitij12345 @janeyx99