added view_as

thunder/executors/torchex.py

@@ -473,6 +473,7 @@ def _tensor_from_sequence_prims_transform(
 unfold = _register_torch_operation("unfold", module=torch.Tensor)
 unsqueeze = _register_torch_operation("unsqueeze")
 view = _register_torch_operation("view", module=torch.Tensor)
+view_as = _register_torch_operation("view_as", module=torch.Tensor)
 
 
 def _broadcast_in_dim_prim_transform(
@@ -565,6 +566,7 @@ def _squeeze_transform(a: TensorLike, /, dim: None | int | Sequence[int] = None)
 _register_implementation(ltorch.unfold, unfold, checker=_always_executable)
 _register_implementation(ltorch.unsqueeze, unsqueeze, checker=_always_executable)
 _register_implementation(ltorch.view, view, checker=_always_executable)
+_register_implementation(ltorch.view_as, view_as, checker=_always_executable)
 
 #
 # Memory format operations

thunder/tests/opinfos.py

@@ -3689,6 +3689,45 @@ def reshape_sample_generator(op, device, dtype, requires_grad, **kwargs):
 shape_ops.append(reshape_opinfo)
 
 
+def view_as_sample_generator(op, device, dtype, requires_grad, **kwargs):
+    make = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad)
+
+    # Input shape, output shape
+    cases = (
+        ((4,), (4,)),  # no-op
+        ((2, 2), (2, 2)),  # no-op
+        ((1, 2, 1), (1, 2, 1)),  # no-op
+        ((2, 2, 2), (4, 2)),
+        ((125,), (25, 5)),
+        ((25, 25), (1, 5, 5, 1, 5, 1, 5, 1)),
+        ((16, 32), (2, 4, 1, 4, 4, 1, 4)),
+        ((16, 12), (12, 16)),
+        ((1, 16, 12), (12, 16)),
+        ((1, 5, 1, 5), (25, 1)),
+        ((2, 4, 2), (4, 4)),
+        ((1, 4), (1, 1, 2, 1, 2)),
+        ((3, 5, 7), (7, 5, 3)),
+        ((1,), ()),  # empty
+        ((5, 0, 2, 3), (5, 0, 2, 3)),
+        ((2, 1, 0, 3, 1), (5, 0)),
+        ((1,), ()),  # empty
+        ((4, 5, 6), (4, 5, 6, 1, 1, 1)),
+        ((), (1, 1, 1, 1)),  # empty
+        ((), ()),
+    )
+
+    for ishape, oshape in cases:
+        yield SampleInput(make(ishape), make(oshape))
+
+
+view_as_opinfo = OpInfo(
+    ltorch.view_as,
+    sample_input_generator=view_as_sample_generator,
+    torch_reference=torch.Tensor.view_as,
+)
+shape_ops.append(view_as_opinfo)
+
+
 def repeat_sample_generator(op, device, dtype, requires_grad, **kwargs):
     make = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad)
 

thunder/torch/__init__.py

@@ -1053,6 +1053,11 @@ def view(a: TensorLike, /, *shape) -> TensorLike:
     return reshape(a, shape)
 
 
+@torchsymbol(torch.Tensor.view_as, is_method=True)
+def view_as(a: TensorLike, b: TensorLike, /) -> TensorLike:
+    return view(a, b.size())
+
+
 #
 # Elementwise unary operaitons
 #