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

ninetoothed/tensor.py

@@ -3,7 +3,6 @@ import math
 import re
 
 import ninetoothed.naming as naming
-from ninetoothed.language import call
 from ninetoothed.symbol import Symbol
 
 
@@ -113,8 +112,11 @@ class Tensor:
             if stride == -1:
                 stride = tile_size
 
+            def cdiv(x, y):
+                return (x + y - 1) // y
+
             new_size = (
-                call("cdiv", self_size - spacing * (tile_size - 1) - 1, stride) + 1
+                (cdiv(self_size - spacing * (tile_size - 1) - 1, stride) + 1)
                 if stride != 0
                 else -1
             )

ninetoothed/visualization.py

@@ -0,0 +1,122 @@
+import matplotlib.pyplot as plt
+import numpy as np
+from mpl_toolkits.axes_grid1 import Divider, Size
+
+
+def visualize(tensor, color=None, save_path=None):
+    outline_width = 0.1
+    plt.rcParams["lines.linewidth"] = 72 * outline_width
+
+    if color is None:
+        color = f"C{visualize.count}"
+
+    _, max_pos_x, max_pos_y = _visualize_tensor(plt.gca(), tensor, 0, 0, color)
+
+    width = max_pos_y + 1
+    height = max_pos_x + 1
+
+    fig = plt.figure(figsize=(width + outline_width, height + outline_width))
+
+    h = (Size.Fixed(0), Size.Fixed(width + outline_width))
+    v = (Size.Fixed(0), Size.Fixed(height + outline_width))
+
+    divider = Divider(fig, (0, 0, 1, 1), h, v, aspect=False)
+
+    ax = fig.add_axes(
+        divider.get_position(), axes_locator=divider.new_locator(nx=1, ny=1)
+    )
+
+    ax.set_aspect("equal")
+    ax.invert_yaxis()
+
+    plt.axis("off")
+
+    half_outline_width = outline_width / 2
+    plt.xlim((-half_outline_width, width + half_outline_width))
+    plt.ylim((-half_outline_width, height + half_outline_width))
+
+    _visualize_tensor(ax, tensor, 0, 0, color)
+
+    plt.savefig(save_path, transparent=True, bbox_inches="tight", pad_inches=0)
+
+    plt.close()
+
+    visualize.count += 1
+
+
+visualize.count = 0
+
+
+def _visualize_tensor(ax, tensor, x, y, color, level_spacing=4):
+    verts = _visualize_level(ax, tensor, x, y, color)
+
+    if tensor.dtype is None:
+        return verts, verts[1][1][0], verts[1][1][1]
+
+    next_x, next_y = verts[0][1]
+    next_y += level_spacing + 1
+
+    next_verts, max_pos_x, max_pos_y = _visualize_tensor(
+        ax, tensor.dtype, next_x, next_y, color
+    )
+
+    conn_verts = _verts_of_rect(1, level_spacing, next_x, next_y - level_spacing)
+    conn_verts = [list(vert) for vert in conn_verts]
+    conn_verts[2][0] += next_verts[1][0][0]
+
+    pos_y, pos_x = zip(*conn_verts)
+    pos_x = pos_x + (pos_x[0],)
+    pos_y = pos_y + (pos_y[0],)
+
+    ax.plot(pos_x[1:3], pos_y[1:3], "k--")
+    ax.plot(pos_x[3:5], pos_y[3:5], "k--")
+
+    max_pos_x = max(max_pos_x, verts[1][1][0])
+    max_pos_y = max(max_pos_y, verts[1][1][1])
+
+    return verts, max_pos_x, max_pos_y
+
+
+def _visualize_level(ax, level, x, y, color):
+    offsets = [1 for _ in range(level.ndim)]
+
+    for dim in range(-3, -level.ndim - 1, -1):
+        offsets[dim] = offsets[dim + 2] * level.shape[dim + 2] + 1
+
+    indices = np.indices(level.shape)
+    flattened_indices = np.stack(
+        [indices[i].flatten() for i in range(level.ndim)], axis=-1
+    )
+
+    max_pos_x = x
+    max_pos_y = y
+
+    for indices in flattened_indices:
+        pos = [x, y]
+
+        for dim, index in enumerate(indices):
+            pos[(level.ndim - dim) % 2] += index * offsets[dim]
+
+        max_pos_x = max(max_pos_x, pos[0])
+        max_pos_y = max(max_pos_y, pos[1])
+
+        _visualize_unit_square(ax, pos[1], pos[0], color)
+
+    verts = (((x, y), (x, max_pos_y)), ((max_pos_x, y), (max_pos_x, max_pos_y)))
+
+    return verts
+
+
+def _visualize_unit_square(ax, x, y, color):
+    _visualize_rect(ax, 1, 1, x, y, color)
+
+
+def _visualize_rect(ax, width, height, x, y, color):
+    pos_x, pos_y = zip(*_verts_of_rect(width, height, x, y))
+
+    ax.fill(pos_x, pos_y, color)
+    ax.plot(pos_x + (pos_x[0],), pos_y + (pos_y[0],), "k")
+
+
+def _verts_of_rect(width, height, x, y):
+    return ((x, y), (x + width, y), (x + width, y + height), (x, y + height))

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: ninetoothed
-Version: 0.11.1
+Version: 0.12.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
@@ -11,6 +11,9 @@ Classifier: Operating System :: OS Independent
 Classifier: Programming Language :: Python :: 3
 Requires-Python: >=3.10
 Requires-Dist: triton>=3.0.0
+Provides-Extra: visualization
+Requires-Dist: matplotlib>=3.9.0; extra == 'visualization'
+Requires-Dist: numpy>=2.1.0; extra == 'visualization'
 Description-Content-Type: text/markdown
 
 # NineToothed

{ninetoothed-0.11.1.dist-info → ninetoothed-0.12.0.dist-info}/RECORD

@@ -3,9 +3,10 @@ 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/tensor.py,sha256=E63sq3jh7ZLiLwFYTtavztKEZx7kRX-UVa2ZXSP2X0s,12008
 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/visualization.py,sha256=VPPh__Bral_Z9hKj9D4UOo8HvRFQidCWSe9cS-D5QfY,3351
+ninetoothed-0.12.0.dist-info/METADATA,sha256=gb5zeAxwYQRm-dFNmdX3uDIn_U-abEZ-VU9bSBBxioA,7198
+ninetoothed-0.12.0.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+ninetoothed-0.12.0.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+ninetoothed-0.12.0.dist-info/RECORD,,