Support tensor.T for transpose (#1110)

Author: yf225

helion/_compiler/type_propagation.py

@@ -444,7 +444,7 @@ def propagate_unary(self, op: ast.unaryop, origin: Origin) -> TypeInfo:
 
     def propagate_attribute(self, attr: str, origin: AttributeOrigin) -> TypeInfo:
         assert origin.key == attr
-        if attr in {"dtype", "device", "ndim", "shape"}:
+        if attr in {"dtype", "device", "ndim", "shape", "T"}:
             return TypeInfo.from_example(getattr(self.fake_value, attr), origin)
         return TensorAttributeType(origin, self)
 

test/test_views.py

@@ -96,6 +96,19 @@ def fn(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
         _code, result = code_and_output(fn, args)
         torch.testing.assert_close(result, args[0] + args[1].transpose(0, 1))
 
+    def test_transpose_T_unsqueeze(self):
+        @helion.kernel(autotune_effort="none")
+        def fn(x: torch.Tensor) -> torch.Tensor:
+            out = torch.empty_like(x)
+            for tile_n, tile_m in hl.tile(x.size()):
+                tile3d = x[tile_n, tile_m].T.unsqueeze(0)
+                out[tile_n, tile_m] = tile3d.squeeze(0).T
+            return out
+
+        args = (torch.randn([512, 384], device=DEVICE),)
+        _, result = code_and_output(fn, args)
+        torch.testing.assert_close(result, args[0])
+
     @unittest.skipUnless(
         supports_tensor_descriptor(), "Tensor descriptor support is required"
     )