[tensor-shapes] Can we just use `*Shape` instead of `*Elements[Shape]`?

[stroxler]

I'm pretty sure that *tuple[A, B] has no meaning today, nor does *T where T: tuple. And similarly for IntTuple.

If that's the case, and we're confident it will continue to be, then Elements might be unnecessary. We can probably just allow the tuple type itself to be splatted directly. That would make code that actually needs such splats quite a bit more compact

Inspired by @TimothyEDawson's observation in #4807

[stroxler]

One pretty clear issue here is that Elements is serving a real purpose at runtime, because unary * is not legal on tuple or TypeVar. So if we do this, users would have to be using from __future__ import annotations.

That might be reason enough not to implement it as just a feature - we might want a typing discussion and to consider a PEP allowing this before we support more notation that breaks at runtime.

Allowing arithmetic ops like + for IntVar is already not runtime safe when using the PEP 695 notation since the IntVars are actually TypeVar at runtime. So we're already flirting with this problem today. But I think the cost-benefit analysis for this sugar is very different than for +

[TimothyEDawson]

One pretty clear issue here is that Elements is serving a real purpose at runtime, because unary * is not legal on tuple or TypeVar. So if we do this, users would have to be using from future import annotations.

For the runtime support, would that not be as simple as adding the __iter__ method to IntTuple? I see what you mean, though.

For whatever it's worth, I came up with a little way to hide *Elements while using the current syntax, given my use case where pretty much everything has a shared initial batch dimension:

from shape_extensions import Elements, IntTuple, IntVar
from torch import Tensor, ones

type S[Batch: IntTuple, Dim: IntVar] = IntTuple[*Elements[Batch], Dim]


def f[Batch: IntTuple, Dim: IntVar](
    x: Tensor[Batch], y: Tensor[S[Batch, Dim]]
) -> Tensor[S[Batch, Dim]]:
    return x[..., None] * y


x = ones((5, 4))
y = ones((5, 4, 3))

z1 = x[..., None] * y  # Tensor[IntTuple[5, 4, 3]]
z2 = f(x, y)  # Tensor[IntTuple[5, 4, 3]]

Issue #4977 was inspired by an attempt to extend this pattern to the case where the inputs have different but broadcastable batch dimensions.

[stroxler]

I'm going ahead with this change

The runtime support issue is mostly around TypeVar, not IntTuple: at runtime, the value of Batch is a TypeVar with bound=IntTuple, and unlike IntTuple we don't support whether TypeVar allows splat.

But it's easy to add support to IntTuple, and after thinking about this I don't think *Batch is very different from N + M - both of these are going to have runtime evaluation issues, and we have workarounds for both of them when needed.