Add squeeze method to Tensor

InfiniTensor · Aug 21, 2024 · ea03213 · ea03213
1 parent 270ba96
commit ea03213
Showing 1 changed file with 9 additions and 0 deletions.
diff --git a/src/ninetoothed/tensor.py b/src/ninetoothed/tensor.py
@@ -93,6 +93,15 @@ def expand(self, shape):
             original=self.original,
         )
 
+    def squeeze(self, dim):
+        # TODO: Add error handling.
+        return type(self)(
+            shape=[size for i, size in enumerate(self.shape) if dim != i],
+            dtype=self.dtype,
+            strides=[stride for i, stride in enumerate(self.strides) if dim != i],
+            original=self.original,
+        )
+
     def names(self):
         return (
             {self.original.pointer_string()}