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

ninetoothed/jit.py

@@ -20,6 +20,13 @@ from ninetoothed.torchifier import Torchifier
 
 
 def make(arrangement, application, tensors):
+    """Integrate the arrangement and the application of the tensors.
+
+    :param arrangement: The arrangement of the tensors.
+    :param application: The application of the tensors.
+    :param tensors: The tensors.
+    :return: A handle to the compute kernel.
+    """
     params = inspect.signature(application).parameters
     types = arrangement(*tensors)
     annotations = {param: type for param, type in zip(params, types)}
@@ -28,14 +35,26 @@ def make(arrangement, application, tensors):
     return jit(application)
 
 
-def jit(_func=None, *, _prettify=False):
+def jit(func=None, *, _prettify=False):
+    """A decorator for generating compute kernels.
+
+    :param func: The function to be compiled.
+    :param _prettify: Whether to prettify the generated code.
+    :return: A handle to the compute kernel.
+
+    .. note::
+
+        The ``_prettify`` parameter is experimental, which might break
+        the generated code.
+    """
+
     def wrapper(func):
         return JIT(func, _prettify=_prettify)()
 
-    if _func is None:
+    if func is None:
         return wrapper
 
-    return wrapper(_func)
+    return wrapper(func)
 
 
 class JIT:
@@ -472,6 +491,15 @@ class CodeGenerator(ast.NodeTransformer):
                 for target_dim in range(tensor.target.ndim)
                 if offsets[source_dim][target_dim] != 0
             ),
+        ) & functools.reduce(
+            lambda x, y: x & y,
+            (
+                indices[dim - tensor.innermost().target.ndim][
+                    type(self)._generate_slices(tensor, target_dim)
+                ]
+                < tensor.innermost().target.shape[dim]
+                for dim, target_dim in enumerate(tensor.innermost().target_dims)
+            ),
         )
 
         return pointers, mask

ninetoothed/symbol.py

@@ -7,6 +7,13 @@ import ninetoothed.naming as naming
 
 
 class Symbol:
+    """A class uesed to represent a symbol.
+
+    :param expr: The expression used to construct the symbol.
+    :param constexpr: Whether the symbol is a constexpr.
+    :param mata: Whether the symbol is a meta.
+    """
+
     def __init__(self, expr, constexpr=None, meta=None):
         if isinstance(expr, type(self)):
             self._node = expr._node

ninetoothed/tensor.py

@@ -8,6 +8,21 @@ from ninetoothed.symbol import Symbol
 
 
 class Tensor:
+    """A class uesed to represent a symbolic tensor.
+
+    :param ndim: The number of dimensions of the tensor.
+    :param shape: The shape of the tensor.
+    :param dtype: The element type of the tensor.
+    :param strides: The strides of the tensor.
+    :param other: The values for out-of-bounds positions.
+    :param constexpr_shape: Whether the sizes are constexpr.
+    :param name: The name of the tensor.
+    :param source: For internal use only.
+    :param source_dims: For internal use only.
+    :param target: For internal use only.
+    :param target_dims: For internal use only.
+    """
+
     num_instances = 0
 
     def __init__(
@@ -70,6 +85,14 @@ class Tensor:
         type(self).num_instances += 1
 
     def tile(self, tile_shape, strides=None, dilation=None):
+        """Tiles the tensor into a hierarchical tensor.
+
+        :param tile_shape: The shape of a tile.
+        :param strides: The interval at which each tile is generated.
+        :param dilation: The spacing between tiles.
+        :return: A hierarchical tensor.
+        """
+
         if strides is None:
             strides = [-1 for _ in tile_shape]
 
@@ -119,6 +142,12 @@ class Tensor:
         )
 
     def expand(self, shape):
+        """Expands the specified singleton dimensions of the tensor.
+
+        :param shape: The expanded shape.
+        :return: The expanded tensor.
+        """
+
         # TODO: Add error handling.
         return type(self)(
             shape=[
@@ -136,6 +165,12 @@ class Tensor:
         )
 
     def squeeze(self, dim):
+        """Removes the specified singleton dimensions of the tensor.
+
+        :param dim: The dimension(s) to be squeezed.
+        :return: The squeezed tensor.
+        """
+
         if not isinstance(dim, tuple):
             dim = (dim,)
 
@@ -158,6 +193,12 @@ class Tensor:
         )
 
     def permute(self, dims):
+        """Permutes the dimensions of the tensor.
+
+        :param dims: The permuted ordering of the dimensions.
+        :return: The permuted tensor.
+        """
+
         # TODO: Add error handling.
         new_shape = [None for _ in range(self.ndim)]
         new_strides = [None for _ in range(self.ndim)]
@@ -178,6 +219,16 @@ class Tensor:
         )
 
     def flatten(self, start_dim=None, end_dim=None):
+        """Flattens the specified dimensions of the tensor.
+
+        See :func:`ravel` for the differences between :func:`flatten`
+        and :func:`ravel`.
+
+        :param start_dim: The first dimension to flatten.
+        :param end_dim: The dimension after the last to flatten.
+        :return: The flattened tensor.
+        """
+
         # TODO: Add error handling.
         if start_dim is None:
             start_dim = 0
@@ -222,6 +273,18 @@ class Tensor:
         )
 
     def ravel(self):
+        """Flattens the hierarchy of the tensor.
+
+        :func:`ravel` differs from :func:`flatten`, which only flattens
+        dimensions at a single level. For example, consider a tensor
+        with two levels: the first level has a shape of ``(N, P, Q)``,
+        and the second level has a shape of ``(C, R, S)``. After
+        applying :func:`ravel`, the resulting tensor will have a single
+        flattened level with a shape of ``(N, P, Q, C, R, S)``.
+
+        :return: The raveled tensor.
+        """
+
         # TODO: Add error handling.
         new_shape = []
         new_strides = []

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: ninetoothed
-Version: 0.11.0
+Version: 0.11.1
 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.1.dist-info/RECORD

@@ -0,0 +1,11 @@
+ninetoothed/__init__.py,sha256=dX34sk5GA3OgWf1Jc4gJMW3UwcGcJsuG3hs3rkiqq6g,161
+ninetoothed/jit.py,sha256=U3Nen5vyx69ulW7_hnRuATW86Ag9NgVgd3U02NVB20c,24430
+ninetoothed/language.py,sha256=YwjlBENmmKPTnhaQ2uYbj5MwzrCAT7MLJ6VkQ6NeXJE,504
+ninetoothed/naming.py,sha256=Fl0x4eDRStTpkXjJg6179ErEnY7bR5Qi0AT6RX9C3fU,951
+ninetoothed/symbol.py,sha256=mN96tp-2eUxbiNfxuxtKWNSxOSdYqlcmpY2MYQ-FiEg,4993
+ninetoothed/tensor.py,sha256=Pgl3t08qNDJrjbNMLltEIwMu19vnKwMUncOq4aedTjY,11983
+ninetoothed/torchifier.py,sha256=aDijK5UOwK2oLXDHgDo8M959rJclEI0lcfaPr7GQTXY,1012
+ninetoothed-0.11.1.dist-info/METADATA,sha256=DsvP5HloDDgEiHUtPEu8blUVYCBr3_VXYILUikfOE-Y,7055
+ninetoothed-0.11.1.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+ninetoothed-0.11.1.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+ninetoothed-0.11.1.dist-info/RECORD,,

ninetoothed-0.11.0.dist-info/RECORD

@@ -1,11 +0,0 @@
-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,,