ninetoothed 0.10.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.10.0 → ninetoothed-0.11.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: ninetoothed
-Version: 0.10.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.10.0 → ninetoothed-0.11.0}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "ninetoothed"
-version = "0.10.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.10.0 → ninetoothed-0.11.0}/src/ninetoothed/tensor.py

@@ -132,6 +132,7 @@ class Tensor:
             ],
             source=self.source,
             source_dims=self.source_dims,
+            target_dims=self.target_dims,
         )
 
     def squeeze(self, dim):
@@ -149,6 +150,11 @@ class Tensor:
                 for i, source_dim in enumerate(self.source_dims)
                 if i not in dim
             ],
+            target_dims=[
+                target_dim
+                for i, target_dim in enumerate(self.target_dims)
+                if i not in dim
+            ],
         )
 
     def permute(self, dims):
@@ -168,6 +174,7 @@ class Tensor:
             strides=new_strides,
             source=self.source,
             source_dims=new_source_dims,
+            target_dims=self.target_dims,
         )
 
     def flatten(self, start_dim=None, end_dim=None):
@@ -197,12 +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_dims=new_target_dims,
         )
 
     def ravel(self):