ninetoothed 0.9.0__tar.gz → 0.11.0__tar.gz

ninetoothed-0.11.0/.github/workflows/publish-to-pypi.yml

@@ -0,0 +1,92 @@
+name: Publish Python 🐍 distribution 📦 to PyPI
+
+on: push
+
+jobs:
+  build:
+    name: Build distribution 📦
+    runs-on: ubuntu-latest
+
+    steps:
+    - uses: actions/checkout@v4
+      with:
+        persist-credentials: false
+    - name: Set up Python
+      uses: actions/setup-python@v5
+      with:
+        python-version: "3.x"
+    - name: Install pypa/build
+      run: >-
+        python3 -m
+        pip install
+        build
+        --user
+    - name: Build a binary wheel and a source tarball
+      run: python3 -m build
+    - name: Store the distribution packages
+      uses: actions/upload-artifact@v4
+      with:
+        name: python-package-distributions
+        path: dist/
+
+  publish-to-pypi:
+    name: >-
+      Publish Python 🐍 distribution 📦 to PyPI
+    if: startsWith(github.ref, 'refs/tags/')
+    needs:
+    - build
+    runs-on: ubuntu-latest
+    environment:
+      name: pypi
+      url: https://pypi.org/p/ninetoothed
+    permissions:
+      id-token: write
+
+    steps:
+    - name: Download all the dists
+      uses: actions/download-artifact@v4
+      with:
+        name: python-package-distributions
+        path: dist/
+    - name: Publish distribution 📦 to PyPI
+      uses: pypa/gh-action-pypi-publish@release/v1
+
+  github-release:
+    name: >-
+      Sign the Python 🐍 distribution 📦 with Sigstore
+      and upload them to GitHub Release
+    needs:
+    - publish-to-pypi
+    runs-on: ubuntu-latest
+
+    permissions:
+      contents: write
+      id-token: write
+
+    steps:
+    - name: Download all the dists
+      uses: actions/download-artifact@v4
+      with:
+        name: python-package-distributions
+        path: dist/
+    - name: Sign the dists with Sigstore
+      uses: sigstore/gh-action-sigstore-python@v3.0.0
+      with:
+        inputs: >-
+          ./dist/*.tar.gz
+          ./dist/*.whl
+    - name: Create GitHub Release
+      env:
+        GITHUB_TOKEN: ${{ github.token }}
+      run: >-
+        gh release create
+        "$GITHUB_REF_NAME"
+        --repo "$GITHUB_REPOSITORY"
+        --notes ""
+    - name: Upload artifact signatures to GitHub Release
+      env:
+        GITHUB_TOKEN: ${{ github.token }}
+      run: >-
+        gh release upload
+        "$GITHUB_REF_NAME" dist/**
+        --repo "$GITHUB_REPOSITORY"

{ninetoothed-0.9.0 → ninetoothed-0.11.0}/PKG-INFO

@@ -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.9.0 → ninetoothed-0.11.0}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "ninetoothed"
-version = "0.9.0"
+version = "0.11.0"
 authors = [{ name = "Jiacheng Huang", email = "huangjiacheng0709@outlook.com" }]
 description = "A domain-specific language based on Triton but providing higher-level abstraction."
 readme = "README.md"

{ninetoothed-0.9.0 → ninetoothed-0.11.0}/src/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-0.9.0 → ninetoothed-0.11.0}/src/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"