Ban some ops

Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

Author: ksivaman

transformer_engine/pytorch/tensor/grouped_tensor.py

@@ -14,6 +14,33 @@
 from .storage.grouped_tensor_storage import GroupedTensorStorage
 
 
+# For now, conservatively ban all shape manipulatimg ops.
+BANNED_SHAPE_OPS = {
+    torch.ops.aten.view.default,
+    torch.ops.aten._unsafe_view.default,
+    torch.ops.aten.reshape.default,
+    torch.ops.aten._reshape_alias.default,
+    torch.ops.aten.flatten.using_ints,
+    torch.ops.aten.unflatten.int,
+    torch.ops.aten.squeeze.dim,
+    torch.ops.aten.squeeze.dims,
+    torch.ops.aten.unsqueeze.default,
+    torch.ops.aten.transpose.int,
+    torch.ops.aten.permute.default,
+    torch.ops.aten.movedim.int,
+    torch.ops.aten.t.default,
+    torch.ops.aten.slice.Tensor,
+    torch.ops.aten.narrow.default,
+    torch.ops.aten.select.int,
+    torch.ops.aten.split.Tensor,
+    torch.ops.aten.chunk.default,
+    torch.ops.aten.expand.default,
+    torch.ops.aten.expand_as.default,
+    torch.ops.aten.cat.default,
+    torch.ops.aten.stack.default,
+}
+
+
 class GroupedTensor(GroupedTensorStorage, torch.Tensor):
     """Tensor wrapper class for grouped tensor storage."""
 
@@ -96,6 +123,10 @@ def __torch_dispatch__(cls, func, types, args, kwargs=None):
         if func in (torch.ops.aten.detach.default, torch.ops.aten.alias.default):
             return args[0]
 
+        # Don't allow reshape/view etc.
+        if func in BANNED_SHAPE_OPS:
+            raise RuntimeError(f"{cls.__name__} forbids shape-manipulation op: {func} ")
+
         def grouped_to_stacked_tensor(grouped: GroupedTensor) -> torch.Tensor:
             if not grouped.all_same_shape():
                 raise NotImplementedError(