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

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
ninetoothed/jit.py CHANGED
@@ -20,6 +20,13 @@ from ninetoothed.torchifier import Torchifier
20
20
 
21
21
 
22
22
  def make(arrangement, application, tensors):
23
+ """Integrate the arrangement and the application of the tensors.
24
+
25
+ :param arrangement: The arrangement of the tensors.
26
+ :param application: The application of the tensors.
27
+ :param tensors: The tensors.
28
+ :return: A handle to the compute kernel.
29
+ """
23
30
  params = inspect.signature(application).parameters
24
31
  types = arrangement(*tensors)
25
32
  annotations = {param: type for param, type in zip(params, types)}
@@ -28,14 +35,26 @@ def make(arrangement, application, tensors):
28
35
  return jit(application)
29
36
 
30
37
 
31
- def jit(_func=None, *, _prettify=False):
38
+ def jit(func=None, *, _prettify=False):
39
+ """A decorator for generating compute kernels.
40
+
41
+ :param func: The function to be compiled.
42
+ :param _prettify: Whether to prettify the generated code.
43
+ :return: A handle to the compute kernel.
44
+
45
+ .. note::
46
+
47
+ The ``_prettify`` parameter is experimental, which might break
48
+ the generated code.
49
+ """
50
+
32
51
  def wrapper(func):
33
52
  return JIT(func, _prettify=_prettify)()
34
53
 
35
- if _func is None:
54
+ if func is None:
36
55
  return wrapper
37
56
 
38
- return wrapper(_func)
57
+ return wrapper(func)
39
58
 
40
59
 
41
60
  class JIT:
@@ -472,6 +491,15 @@ class CodeGenerator(ast.NodeTransformer):
472
491
  for target_dim in range(tensor.target.ndim)
473
492
  if offsets[source_dim][target_dim] != 0
474
493
  ),
494
+ ) & functools.reduce(
495
+ lambda x, y: x & y,
496
+ (
497
+ indices[dim - tensor.innermost().target.ndim][
498
+ type(self)._generate_slices(tensor, target_dim)
499
+ ]
500
+ < tensor.innermost().target.shape[dim]
501
+ for dim, target_dim in enumerate(tensor.innermost().target_dims)
502
+ ),
475
503
  )
476
504
 
477
505
  return pointers, mask
ninetoothed/symbol.py CHANGED
@@ -7,6 +7,13 @@ import ninetoothed.naming as naming
7
7
 
8
8
 
9
9
  class Symbol:
10
+ """A class uesed to represent a symbol.
11
+
12
+ :param expr: The expression used to construct the symbol.
13
+ :param constexpr: Whether the symbol is a constexpr.
14
+ :param mata: Whether the symbol is a meta.
15
+ """
16
+
10
17
  def __init__(self, expr, constexpr=None, meta=None):
11
18
  if isinstance(expr, type(self)):
12
19
  self._node = expr._node
ninetoothed/tensor.py CHANGED
@@ -8,6 +8,21 @@ from ninetoothed.symbol import Symbol
8
8
 
9
9
 
10
10
  class Tensor:
11
+ """A class uesed to represent a symbolic tensor.
12
+
13
+ :param ndim: The number of dimensions of the tensor.
14
+ :param shape: The shape of the tensor.
15
+ :param dtype: The element type of the tensor.
16
+ :param strides: The strides of the tensor.
17
+ :param other: The values for out-of-bounds positions.
18
+ :param constexpr_shape: Whether the sizes are constexpr.
19
+ :param name: The name of the tensor.
20
+ :param source: For internal use only.
21
+ :param source_dims: For internal use only.
22
+ :param target: For internal use only.
23
+ :param target_dims: For internal use only.
24
+ """
25
+
11
26
  num_instances = 0
12
27
 
13
28
  def __init__(
@@ -70,6 +85,14 @@ class Tensor:
70
85
  type(self).num_instances += 1
71
86
 
72
87
  def tile(self, tile_shape, strides=None, dilation=None):
88
+ """Tiles the tensor into a hierarchical tensor.
89
+
90
+ :param tile_shape: The shape of a tile.
91
+ :param strides: The interval at which each tile is generated.
92
+ :param dilation: The spacing between tiles.
93
+ :return: A hierarchical tensor.
94
+ """
95
+
73
96
  if strides is None:
74
97
  strides = [-1 for _ in tile_shape]
75
98
 
@@ -119,6 +142,12 @@ class Tensor:
119
142
  )
120
143
 
121
144
  def expand(self, shape):
145
+ """Expands the specified singleton dimensions of the tensor.
146
+
147
+ :param shape: The expanded shape.
148
+ :return: The expanded tensor.
149
+ """
150
+
122
151
  # TODO: Add error handling.
123
152
  return type(self)(
124
153
  shape=[
@@ -136,6 +165,12 @@ class Tensor:
136
165
  )
137
166
 
138
167
  def squeeze(self, dim):
168
+ """Removes the specified singleton dimensions of the tensor.
169
+
170
+ :param dim: The dimension(s) to be squeezed.
171
+ :return: The squeezed tensor.
172
+ """
173
+
139
174
  if not isinstance(dim, tuple):
140
175
  dim = (dim,)
141
176
 
@@ -158,6 +193,12 @@ class Tensor:
158
193
  )
159
194
 
160
195
  def permute(self, dims):
196
+ """Permutes the dimensions of the tensor.
197
+
198
+ :param dims: The permuted ordering of the dimensions.
199
+ :return: The permuted tensor.
200
+ """
201
+
161
202
  # TODO: Add error handling.
162
203
  new_shape = [None for _ in range(self.ndim)]
163
204
  new_strides = [None for _ in range(self.ndim)]
@@ -178,6 +219,16 @@ class Tensor:
178
219
  )
179
220
 
180
221
  def flatten(self, start_dim=None, end_dim=None):
222
+ """Flattens the specified dimensions of the tensor.
223
+
224
+ See :func:`ravel` for the differences between :func:`flatten`
225
+ and :func:`ravel`.
226
+
227
+ :param start_dim: The first dimension to flatten.
228
+ :param end_dim: The dimension after the last to flatten.
229
+ :return: The flattened tensor.
230
+ """
231
+
181
232
  # TODO: Add error handling.
182
233
  if start_dim is None:
183
234
  start_dim = 0
@@ -222,6 +273,18 @@ class Tensor:
222
273
  )
223
274
 
224
275
  def ravel(self):
276
+ """Flattens the hierarchy of the tensor.
277
+
278
+ :func:`ravel` differs from :func:`flatten`, which only flattens
279
+ dimensions at a single level. For example, consider a tensor
280
+ with two levels: the first level has a shape of ``(N, P, Q)``,
281
+ and the second level has a shape of ``(C, R, S)``. After
282
+ applying :func:`ravel`, the resulting tensor will have a single
283
+ flattened level with a shape of ``(N, P, Q, C, R, S)``.
284
+
285
+ :return: The raveled tensor.
286
+ """
287
+
225
288
  # TODO: Add error handling.
226
289
  new_shape = []
227
290
  new_strides = []
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: ninetoothed
3
- Version: 0.11.0
3
+ Version: 0.11.1
4
4
  Summary: A domain-specific language based on Triton but providing higher-level abstraction.
5
5
  Project-URL: Homepage, https://github.com/InfiniTensor/ninetoothed
6
6
  Project-URL: Issues, https://github.com/InfiniTensor/ninetoothed/issues
@@ -0,0 +1,11 @@
1
+ ninetoothed/__init__.py,sha256=dX34sk5GA3OgWf1Jc4gJMW3UwcGcJsuG3hs3rkiqq6g,161
2
+ ninetoothed/jit.py,sha256=U3Nen5vyx69ulW7_hnRuATW86Ag9NgVgd3U02NVB20c,24430
3
+ ninetoothed/language.py,sha256=YwjlBENmmKPTnhaQ2uYbj5MwzrCAT7MLJ6VkQ6NeXJE,504
4
+ ninetoothed/naming.py,sha256=Fl0x4eDRStTpkXjJg6179ErEnY7bR5Qi0AT6RX9C3fU,951
5
+ ninetoothed/symbol.py,sha256=mN96tp-2eUxbiNfxuxtKWNSxOSdYqlcmpY2MYQ-FiEg,4993
6
+ ninetoothed/tensor.py,sha256=Pgl3t08qNDJrjbNMLltEIwMu19vnKwMUncOq4aedTjY,11983
7
+ ninetoothed/torchifier.py,sha256=aDijK5UOwK2oLXDHgDo8M959rJclEI0lcfaPr7GQTXY,1012
8
+ ninetoothed-0.11.1.dist-info/METADATA,sha256=DsvP5HloDDgEiHUtPEu8blUVYCBr3_VXYILUikfOE-Y,7055
9
+ ninetoothed-0.11.1.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
10
+ ninetoothed-0.11.1.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
11
+ ninetoothed-0.11.1.dist-info/RECORD,,
@@ -1,11 +0,0 @@
1
- ninetoothed/__init__.py,sha256=dX34sk5GA3OgWf1Jc4gJMW3UwcGcJsuG3hs3rkiqq6g,161
2
- ninetoothed/jit.py,sha256=0LeDBpSYFgPx4hatP_ZsvElsj0d9d552OKRc__L1Jvc,23460
3
- ninetoothed/language.py,sha256=YwjlBENmmKPTnhaQ2uYbj5MwzrCAT7MLJ6VkQ6NeXJE,504
4
- ninetoothed/naming.py,sha256=Fl0x4eDRStTpkXjJg6179ErEnY7bR5Qi0AT6RX9C3fU,951
5
- ninetoothed/symbol.py,sha256=rZ5nXtn-U1Nw0BBRJ-kfrwmX_zCbAi76un-Z2QFaoZc,4773
6
- ninetoothed/tensor.py,sha256=ehgbHrxL1gc9iOI9rGUp5k-nI_daLmCfOdLI9hE-GLw,9756
7
- ninetoothed/torchifier.py,sha256=aDijK5UOwK2oLXDHgDo8M959rJclEI0lcfaPr7GQTXY,1012
8
- ninetoothed-0.11.0.dist-info/METADATA,sha256=NlC4oj1R7gNoCdA1AxYc_7UWY3uZPv3LeOq9Jo-9Q8w,7055
9
- ninetoothed-0.11.0.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
10
- ninetoothed-0.11.0.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
11
- ninetoothed-0.11.0.dist-info/RECORD,,