[tensor-shapes] Combination of Elements and DSL functions

[TimothyEDawson]

Describe the Bug

@stroxler It seems that Elements is not able to unpack the result of a DSL function. While I could write a custom DSL function to handle such shape transformations, I think it's worth noting that a user might expect to be able to do something like this:

from typing import reveal_type

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


def f[B1: IntTuple, B2: IntTuple, Dim: IntVar](
    x: Tensor[B1], y: Tensor[[*Elements[B2], Dim]]
) -> Tensor[[*Elements[broadcast(B1, B2)], Dim]]:
    return x[..., None] * y


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

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

Result of pyrefly check:

ERROR Function call cannot be used in annotations [invalid-annotation]
 --> test2.py:9:24
  |
9 | ) -> Tensor[[*Elements[broadcast(B1, B2)], Dim]]:
  |                        ^^^^^^^^^^^^^^^^^
ERROR `Elements[...]` requires an `IntTuple` or integer tuple, got `Unknown` [invalid-annotation]
 --> test2.py:9:24
  |
9 | ) -> Tensor[[*Elements[broadcast(B1, B2)], Dim]]:
  |                        ^^^^^^^^^^^^^^^^^
 INFO revealed type: Tensor[IntTuple[5, 4, 3]] [reveal-type]
  --> test2.py:17:12
   |
17 | reveal_type(z1)
   |            ----
 INFO revealed type: Tensor[Unknown] [reveal-type]
  --> test2.py:19:12
   |
19 | reveal_type(z2)
   |            ----
 INFO 2 errors

Sandbox Link

No response

(Only applicable for extension issues) IDE Information

No response

[stroxler]

Good catch, I don't think this was intended - the DSL function is conceptually supposed to be legal wherever it's output domain (Int or IntTuple) is.

We might need to audit the rules. Note that the example is an implementation marked with a DSL transform, which at the moment we don't expect to work. But the problem where broadcast should be able to expand through Elements is separate, and I think it's just a bug

[TimothyEDawson]

Right, the error is exactly the same if you had the same signature as a stub.