ninetoothed 0.10.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=[
@@ -132,9 +161,16 @@ class Tensor:
132
161
  ],
133
162
  source=self.source,
134
163
  source_dims=self.source_dims,
164
+ target_dims=self.target_dims,
135
165
  )
136
166
 
137
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
+
138
174
  if not isinstance(dim, tuple):
139
175
  dim = (dim,)
140
176
 
@@ -149,9 +185,20 @@ class Tensor:
149
185
  for i, source_dim in enumerate(self.source_dims)
150
186
  if i not in dim
151
187
  ],
188
+ target_dims=[
189
+ target_dim
190
+ for i, target_dim in enumerate(self.target_dims)
191
+ if i not in dim
192
+ ],
152
193
  )
153
194
 
154
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
+
155
202
  # TODO: Add error handling.
156
203
  new_shape = [None for _ in range(self.ndim)]
157
204
  new_strides = [None for _ in range(self.ndim)]
@@ -168,9 +215,20 @@ class Tensor:
168
215
  strides=new_strides,
169
216
  source=self.source,
170
217
  source_dims=new_source_dims,
218
+ target_dims=self.target_dims,
171
219
  )
172
220
 
173
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
+
174
232
  # TODO: Add error handling.
175
233
  if start_dim is None:
176
234
  start_dim = 0
@@ -197,15 +255,36 @@ class Tensor:
197
255
  leading_source_dims + (flattening_source_dims,) + trailing_source_dims
198
256
  )
199
257
 
258
+ leading_target_dims = self.target_dims[:start_dim]
259
+ flattening_target_dims = self.target_dims[start_dim:end_dim]
260
+ trailing_target_dims = self.target_dims[end_dim:]
261
+
262
+ new_target_dims = (
263
+ leading_target_dims + (flattening_target_dims[-1],) + trailing_target_dims
264
+ )
265
+
200
266
  return type(self)(
201
267
  shape=new_shape,
202
268
  dtype=self.dtype,
203
269
  strides=new_strides,
204
270
  source=self.source,
205
271
  source_dims=new_source_dims,
272
+ target_dims=new_target_dims,
206
273
  )
207
274
 
208
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
+
209
288
  # TODO: Add error handling.
210
289
  new_shape = []
211
290
  new_strides = []
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: ninetoothed
3
- Version: 0.10.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=LGS9wYmPKckZjIEXsMHdclTmVhzkfrc3avOPiCQY1tU,9153
7
- ninetoothed/torchifier.py,sha256=aDijK5UOwK2oLXDHgDo8M959rJclEI0lcfaPr7GQTXY,1012
8
- ninetoothed-0.10.0.dist-info/METADATA,sha256=nQWkQ--AceNN3DoD-kP9_aykZYVw8LOfqf2iO63v1Ek,7055
9
- ninetoothed-0.10.0.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
10
- ninetoothed-0.10.0.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
11
- ninetoothed-0.10.0.dist-info/RECORD,,