ninetoothed 0.9.0__py3-none-any.whl → 0.11.0__py3-none-any.whl

ninetoothed/jit.py

@@ -479,7 +479,7 @@ class CodeGenerator(ast.NodeTransformer):
     def _complete_indices(self, tensor, indices):
         indices = list(self._generate_pid_indices(tensor) + tuple(indices))
 
-        for size in tensor.inmost().shape:
+        for size in tensor.innermost().shape:
             if Symbol.is_name(size):
                 name = size.node.id
                 if not naming.is_meta(name):
@@ -514,7 +514,10 @@ class CodeGenerator(ast.NodeTransformer):
 
     @staticmethod
     def _generate_slices(tensor, dim):
-        return tuple(slice(None) if i == dim else None for i in range(tensor.ndim))
+        return tuple(
+            slice(None) if target_dim == dim else None
+            for target_dim in tensor.innermost().target_dims
+        )
 
     @staticmethod
     def _generate_offsets(tensor, indices):

ninetoothed/tensor.py

@@ -112,14 +112,10 @@ class Tensor:
                 strides=inner_strides,
                 source=self.source,
                 source_dims=self.source_dims,
-                target=self.target,
-                target_dims=self.target_dims,
             ),
             strides=outer_strides,
             source=self.source,
             source_dims=self.source_dims,
-            target=self.target,
-            target_dims=self.target_dims,
         )
 
     def expand(self, shape):
@@ -136,23 +132,28 @@ class Tensor:
             ],
             source=self.source,
             source_dims=self.source_dims,
-            target=self.target,
             target_dims=self.target_dims,
         )
 
     def squeeze(self, dim):
+        if not isinstance(dim, tuple):
+            dim = (dim,)
+
         # TODO: Add error handling.
         return type(self)(
-            shape=[size for i, size in enumerate(self.shape) if dim != i],
+            shape=[size for i, size in enumerate(self.shape) if i not in dim],
             dtype=self.dtype,
-            strides=[stride for i, stride in enumerate(self.strides) if dim != i],
+            strides=[stride for i, stride in enumerate(self.strides) if i not in dim],
             source=self.source,
             source_dims=[
-                source_dim for i, source_dim in enumerate(self.source_dims) if dim != i
+                source_dim
+                for i, source_dim in enumerate(self.source_dims)
+                if i not in dim
             ],
-            target=self.target,
             target_dims=[
-                target_dim for i, target_dim in enumerate(self.target_dims) if dim != i
+                target_dim
+                for i, target_dim in enumerate(self.target_dims)
+                if i not in dim
             ],
         )
 
@@ -173,7 +174,6 @@ class Tensor:
             strides=new_strides,
             source=self.source,
             source_dims=new_source_dims,
-            target=self.target,
             target_dims=self.target_dims,
         )
 
@@ -204,14 +204,21 @@ class Tensor:
             leading_source_dims + (flattening_source_dims,) + trailing_source_dims
         )
 
+        leading_target_dims = self.target_dims[:start_dim]
+        flattening_target_dims = self.target_dims[start_dim:end_dim]
+        trailing_target_dims = self.target_dims[end_dim:]
+
+        new_target_dims = (
+            leading_target_dims + (flattening_target_dims[-1],) + trailing_target_dims
+        )
+
         return type(self)(
             shape=new_shape,
             dtype=self.dtype,
             strides=new_strides,
             source=self.source,
             source_dims=new_source_dims,
-            target=self.target,
-            target_dims=self.target_dims,
+            target_dims=new_target_dims,
         )
 
     def ravel(self):
@@ -250,11 +257,11 @@ class Tensor:
             | (self.source.names() if self.source is not self else set())
         )
 
-    def inmost(self):
+    def innermost(self):
         if not isinstance(self.dtype, type(self)):
             return self
 
-        return self.dtype.inmost()
+        return self.dtype.innermost()
 
     def pointer_string(self):
         return f"{self.name}_pointer"

{ninetoothed-0.9.0.dist-info → ninetoothed-0.11.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: ninetoothed
-Version: 0.9.0
+Version: 0.11.0
 Summary: A domain-specific language based on Triton but providing higher-level abstraction.
 Project-URL: Homepage, https://github.com/InfiniTensor/ninetoothed
 Project-URL: Issues, https://github.com/InfiniTensor/ninetoothed/issues

ninetoothed-0.11.0.dist-info/RECORD

@@ -0,0 +1,11 @@
+ninetoothed/__init__.py,sha256=dX34sk5GA3OgWf1Jc4gJMW3UwcGcJsuG3hs3rkiqq6g,161
+ninetoothed/jit.py,sha256=0LeDBpSYFgPx4hatP_ZsvElsj0d9d552OKRc__L1Jvc,23460
+ninetoothed/language.py,sha256=YwjlBENmmKPTnhaQ2uYbj5MwzrCAT7MLJ6VkQ6NeXJE,504
+ninetoothed/naming.py,sha256=Fl0x4eDRStTpkXjJg6179ErEnY7bR5Qi0AT6RX9C3fU,951
+ninetoothed/symbol.py,sha256=rZ5nXtn-U1Nw0BBRJ-kfrwmX_zCbAi76un-Z2QFaoZc,4773
+ninetoothed/tensor.py,sha256=ehgbHrxL1gc9iOI9rGUp5k-nI_daLmCfOdLI9hE-GLw,9756
+ninetoothed/torchifier.py,sha256=aDijK5UOwK2oLXDHgDo8M959rJclEI0lcfaPr7GQTXY,1012
+ninetoothed-0.11.0.dist-info/METADATA,sha256=NlC4oj1R7gNoCdA1AxYc_7UWY3uZPv3LeOq9Jo-9Q8w,7055
+ninetoothed-0.11.0.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+ninetoothed-0.11.0.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+ninetoothed-0.11.0.dist-info/RECORD,,

ninetoothed-0.9.0.dist-info/RECORD

@@ -1,11 +0,0 @@
-ninetoothed/__init__.py,sha256=dX34sk5GA3OgWf1Jc4gJMW3UwcGcJsuG3hs3rkiqq6g,161
-ninetoothed/jit.py,sha256=B3q32ksKTRr7I4jLcoDXjwEx7A_Awz9DXGEmIkrtoBc,23393
-ninetoothed/language.py,sha256=YwjlBENmmKPTnhaQ2uYbj5MwzrCAT7MLJ6VkQ6NeXJE,504
-ninetoothed/naming.py,sha256=Fl0x4eDRStTpkXjJg6179ErEnY7bR5Qi0AT6RX9C3fU,951
-ninetoothed/symbol.py,sha256=rZ5nXtn-U1Nw0BBRJ-kfrwmX_zCbAi76un-Z2QFaoZc,4773
-ninetoothed/tensor.py,sha256=OU6lVjzKU614mk3EN1AAgTDradbvJkyl22AXwdhxcfs,9577
-ninetoothed/torchifier.py,sha256=aDijK5UOwK2oLXDHgDo8M959rJclEI0lcfaPr7GQTXY,1012
-ninetoothed-0.9.0.dist-info/METADATA,sha256=LKtSbgc_mWKJ4L_8BYiRKWjoV-JKjr4hefhrkdPmrHs,7054
-ninetoothed-0.9.0.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-ninetoothed-0.9.0.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-ninetoothed-0.9.0.dist-info/RECORD,,