PyCuTe ships several printing and visualization utilities that, together, cover almost all of your debugging and explanation needs:
print(layout)/print(tensor)— the canonical(shape):(stride)notation (built into every PyCuTe object).print_tensor— render any rank-1 through rank-4 layout or tensor as nested ASCII tables.print_table— render a rank-2 layout as a single bordered grid (usestabulate).draw_svg— render a rank-2 layout as a colored SVG file: a grid of coordinates, each labelled with the value the layout sends it to.draw_svg_inv— the inverse view: render a layout that sends values to the coordinates of a grid, each coordinate labelled with the value sent there.draw_svg_tv—draw_svg_invof a rank-2 thread-value layout, withT#/V#annotations.draw_latex— render a rank-2 layout as a TikZ/LaTeX document (and PDF); the analogue ofcute::print_latex.draw_latex_invanddraw_latex_tv— the TikZ/LaTeX analogues ofdraw_svg_invanddraw_svg_tv.draw_colors— a small catalog of coloring functors (each returns an(r, g, b)RGB-255 tuple) shared by the SVG and LaTeX drawers:white(the default fordraw_svg/draw_latex),thread_color_8x(tid, vid)(the default for the thread-value drawers),index_grey_8x(idx), and thebank_color_*palettes. See Coloring functors.draw_labels— the labelling functors (each returns a string) of the same drawers:tv_label(tid, vid)(theT#/V#cell label of the thread-value drawers) andindex_label(axis, i)(the row and column indices of every drawer). See Labelling functors.
Every utility lives in pycute.util, which is a
separate import from the core algebra:
>>> from pycute import *
>>> from pycute.util import print_tensor, print_table
>>> from pycute.util import draw_svg, draw_svg_inv, draw_svg_tv
>>> from pycute.util import draw_latex, draw_latex_inv, draw_latex_tv
>>> from pycute.util import white, thread_color_8x # default color functors
>>> from pycute.util import index_grey_8x, bank_color_32x # ...and some alternatives
>>> from pycute.util import tv_label, index_label # default label functorsThe split is intentional: the core algebra has no third-party dependencies.
print_table requires tabulate and
the SVG helpers (draw_svg, draw_svg_inv, draw_svg_tv) require
svgwrite. The LaTeX helpers
(draw_latex, draw_latex_inv, draw_latex_tv) need no Python packages at all —
only a LaTeX installation (e.g. TeX Live providing pdflatex) for the
optional PDF step.
Every PyCuTe Layout, Tensor, ArithTuple, F2, and Swizzle
defines __str__ and __repr__:
>>> from pycute import *
>>> A = Layout((3, (2, 4)), (2, (1, 6)))
>>> print(A)
(3, (2, 4)):(2, (1, 6))
>>> repr(A)
'Layout((3, (2, 4)), (2, (1, 6)))'
>>> T = make_tensor(Layout((4, 4), (4, 1)))
>>> print(T)
<...Array...> o (4, 4):(4, 1) # accessor o layout__str__ is the human-readable shape-and-stride notation. __repr__ is
the full constructor form so you can copy-paste it into an interpreter.
print_tensor(t, print_type=True) is defined in
print_tensor.py. It walks a rank-1
through rank-4 layout or tensor and prints it as nested ASCII tables.
>>> from pycute.util import print_tensor
>>> print_tensor(Layout(8, 1))
8:1
0 1 2 3 4 5 6 7>>> print_tensor(Layout((4, 8), (1, 4)))
(4, 8):(1, 4)
0 4 8 12 16 20 24 28
1 5 9 13 17 21 25 29
2 6 10 14 18 22 26 30
3 7 11 15 19 23 27 31
>>> print_tensor(Layout((3, (2, 3))))
(3, (2, 3)):(1, (3, 6))
0 3 6 9 12 15
1 4 7 10 13 16
2 5 8 11 14 17The header line is suppressed by passing print_type=False.
Higher-rank tensors are printed as 2-D slabs, separated by axis labels.
For a rank-3 tensor, you see one slab per k (the first slab, k = 0,
is printed without a separator line):
[ ... rank-2 slab for k=0 ... ]
-------- k = 1 ---------
[ ... rank-2 slab for k=1 ... ]
For a rank-4 tensor, slabs are grouped by the outermost p and separated
by = lines:
[ ... rank-3 slab for p=0 ... ]
========== p = 1 ===========
[ ... rank-3 slab for p=1 ... ]
This is the same convention as cute::print_tensor in C++ CuTe.
When called on a Layout, print_tensor constructs a tensor backed by an
ImplicitAccessor and prints offsets:
>>> print_tensor(Layout((4, 8), (1, 4)))
(4, 8):(1, 4)
0 4 8 12 16 20 24 28
1 5 9 13 17 21 25 29
2 6 10 14 18 22 26 30
3 7 11 15 19 23 27 31This is the analogue of cute::print_layout in C++ CuTe. It is the
default visualization in this documentation.
print_table(t, print_type=True) renders a rank-2 layout or tensor as a
single bordered grid (one cell per coordinate, the offset inside it). Like
print_tensor, it accepts either a Tensor or a Layout (the latter is
rendered through an ImplicitAccessor) and prints the Shape:Stride
header first unless print_type=False. It uses
tabulate for the table formatting,
so the output is friendlier than print_tensor when you want
copy-pasteable Markdown or fixed-width text:
>>> from pycute.util import print_table
>>> print_table(Layout((4, 8), (1, 4)))
(4, 8):(1, 4)
+---+---+----+----+----+----+----+----+
| 0 | 4 | 8 | 12 | 16 | 20 | 24 | 28 |
+---+---+----+----+----+----+----+----+
| 1 | 5 | 9 | 13 | 17 | 21 | 25 | 29 |
+---+---+----+----+----+----+----+----+
| 2 | 6 | 10 | 14 | 18 | 22 | 26 | 30 |
+---+---+----+----+----+----+----+----+
| 3 | 7 | 11 | 15 | 19 | 23 | 27 | 31 |
+---+---+----+----+----+----+----+----+For non-rank-2 inputs, print_table falls back to a plain print.
draw_svg(tensor, filename="layout.svg", color=white, cell_width=20, cell_label=str, axis_label=index_label)
visualizes crd -> val: it produces an SVG file rendering a layout as a grid
of coordinates crd, each cell labeled with cell_label(val) for the value
val there, its offset. Row and column indices are labeled along the left
and top edges.
Cells are white by default; pass a color functor to shade them, and a
cell_label or axis_label functor to relabel them.
>>> from pycute.util import draw_svg
>>> draw_svg(Layout((4, 8), (1, 4)))
Saved as layout.svg
>>> draw_svg(Layout(((2, 2), (4, 2)), ((1, 8), (2, 16))),
... filename="mixed.svg")
Saved as mixed.svgLike print_tensor, it accepts a Tensor as well as a
Layout. A Tensor labels each cell with the element stored there while
the color still keys off the offset, so one figure shows both the logical
contents and where they live:
>>> T = make_tensor(Layout((4, 8), (1, 4)))
>>> draw_svg(T, filename="data.svg", color=index_grey_8x)
Saved as data.svgThis renders an SVG grid — useful for slides, papers, and quick visual
inspection. For LaTeX/TikZ output, use draw_latex below.
The arguments:
tensor— a rank-2TensororLayout;size[0]rows andsize[1]columns are drawn. A rank-1 input is drawn as a single row.filename— defaults to"layout.svg".color— a coloring functorcolor(idx) -> (r, g, b)mapping a cell's offset to an RGB-255 tuple (each component in[0, 255]). Defaults towhite; the catalog in Coloring functors covers greyscale and bank palettes. The functors coerce the offset withint(), soF2-strided (swizzled) layouts color correctly too.cell_width— the width of a cell in pixels,20by default. Widen it for labels longer than a few characters, such as coordinates.cell_label— a labelling functorcell_label(value) -> strwriting a cell's value: its offset, or the element of aTensor. Defaults tostr.axis_label— a labelling functoraxis_label(axis, i) -> strwriting the indexiof a row (axis == 0) or a column (axis == 1). Defaults toindex_label, the index itself. See Labelling functors for both.
The cell height and fonts are fixed — index labels are monospace, cell
labels are Arial; the palette is controlled by color, the text by
cell_label and axis_label.
For example, to highlight even versus odd offsets:
>>> draw_svg(Layout((4, 8), (1, 4)),
... color=lambda idx: (175, 255, 175) if idx % 2 == 0 else (255, 175, 175))A coordinate prints in basis form when only one axis is nonzero, 7@1 for
(0, 7), so to label every cell of an identity_tensor alike, write its
coordinates with a cell_label:
>>> draw_svg(identity_tensor((4, 8)), filename="coords.svg", cell_width=32,
... cell_label=lambda crd: "({},{})".format(*as_tuple(crd)))
Saved as coords.svgdraw_svg_inv(layout, tile_mn=None, filename="invlayout.svg", color=white, cell_label=str, axis_label=index_label)
visualizes the inverse, val -> crd: composition(tile_mn, layout) sends
each index val of the layout's domain to a coordinate crd of the tile
tile_mn, and that cell is labeled with cell_label(val) and colored by
color(val). The composition folds a linear offset into the tile's (m, n)
coordinates, and leaves (m, n) coordinates — a layout with E(0) / E(1)
strides — as they are. Where several values land on one cell the least is
drawn, and a cell no value lands on is left white.
The two views are inverse to each other: drawn with draw_svg_inv, the
inverse of a compact layout L labels every coordinate with the offset
draw_svg(L) gives it. For a hierarchical L, fold it into the flat tile
product_each(shape(L)).
>>> from pycute.util import draw_svg_inv
>>> L = Layout((4, 8), (8, 1)) # crd -> offset
>>> draw_svg(L, filename="layout.svg")
Saved as layout.svg
>>> draw_svg_inv(right_inverse(L), shape(L), filename="inverse.svg") # offset -> crd
Saved as inverse.svgThe arguments:
layout— a layout, of any rank, whose codomain is a linear offset or(m, n)coordinates.tile_mn— the(M, N)tile drawn. Defaults tocoshape(layout), the extent of(m, n)coordinates; a linear offset needs it given.filename— defaults to"invlayout.svg".color— a coloring functorcolor(val) -> (r, g, b), keyed on the indexvalasdraw_svg's is on the offset, so the same catalog applies. Defaults towhite.cell_label— a labelling functorcell_label(val) -> str. Defaults tostr.axis_label— as indraw_svg.
Errors raised:
Expected a rank-2 MN Tile—tile_mnis not a rank-2 shape of integers. You also get this when you pass a linear-offset layout withouttile_mn, since the defaultcoshape(layout)is then a scalar.
draw_svg_tv(layout, tile_mn=None, filename="tvlayout.svg", color=thread_color_8x, cell_label=tv_label, axis_label=index_label)
is draw_svg_inv for a thread-value layout: a rank-2 layout whose
domain is (tid, vid). This is exactly the layout that comes out of
partitioning a tile across the threads of a warp/threadblock. Each index of
the domain is taken apart into its thread and value, so the functors take
(tid, vid), and the layout is composed into the tile as in draw_svg_inv.
>>> from pycute.util import draw_svg_tv
>>> # (a) codomain already in (m, n) coordinates:
>>> tv = Layout((8, 8), (E(0), E(1)))
>>> draw_svg_tv(tv, filename="identity_tv.svg")
Saved as identity_tv.svg
>>> # (b) codomain as linear offsets: pass the (M, N) tile shape and the
>>> # offsets are mapped into it automatically:
>>> SM80 = Layout(((4, 8), (2, 2)), ((32, 1), (16, 8))) # (tid, vid) -> offset
>>> draw_svg_tv(SM80, tile_mn=(16, 8), filename="sm80_tv.svg")
Saved as sm80_tv.svgEach cell of the resulting SVG shows (by default) T<tid> / V<vid> for
the (tid, vid) pair that lands on that (m, n) position, with the cell
colored (by default) by tid mod 8. Where several pairs land on one cell,
the one of least index tid + T*vid in the domain is drawn, T the number
of threads. Empty positions (no thread covers them) are left white.
The arguments are those of draw_svg_inv, except:
layout— a rank-2 layout with domain(tid, vid).filename— defaults to"tvlayout.svg".color— a coloring functorcolor(tid, vid) -> (r, g, b)mapping a(tid, vid)pair to an RGB-255 tuple (each component in[0, 255]). Defaults tothread_color_8x(coloring bytid % 8), importable frompycute.util. Because it receives bothtidandvid, you can color by value index instead (e.g.color=lambda tid, vid: (200, 200, 255) if vid % 2 else (255, 200, 200)).cell_label— a labelling functorcell_label(tid, vid) -> strwriting the cell of a(tid, vid)pair. Defaults totv_label,T<tid>overV<vid>.
Errors raised: those of draw_svg_inv, and
Expected a rank-2 TV Layout— the layout's domain is not rank 2.
draw_latex(tensor, filename="layout.tex", compile_pdf=True, color=white, cell_label=str, axis_label=index_label)
is the LaTeX/PDF analogue of draw_svg. It writes a standalone
TikZ document rendering a rank-2 layout as a grid — cells
labeled by offset and (by default) unshaded, with row/column index labels —
and then compiles it to a cropped PDF with pdflatex. The output mirrors
cute::print_latex in C++ CuTe, so it is ideal for slides and papers.
>>> from pycute.util import draw_latex
>>> draw_latex(Layout((4, 8), (1, 4)))
Saved as layout.tex
Saved as layout.pdfThe arguments:
tensor— a rank-2TensororLayout. As withdraw_svg, aTensorlabels each cell with the element stored there rather than its offset, and a rank-1 input is drawn as a single row.filename— the.texpath; defaults to"layout.tex". The PDF is written next to it with the same stem (layout.pdf).compile_pdf— whenTrue(default), runpdflatexto produce the PDF and remove the.aux/.logbyproducts. WhenFalse, only the.texis written.color— a coloring functorcolor(idx) -> (r, g, b), identical in contract todraw_svg'scolor(RGB-255 tuple, components in[0, 255]). Defaults towhite.cell_label,axis_label— labelling functors, identical in contract todraw_svg's. A label of several lines is set with\shortstack.
Unlike draw_svg, draw_latex needs no Python packages — only a LaTeX
installation (e.g. TeX Live providing pdflatex) for the PDF step. If
pdflatex is not found on PATH, the .tex is still written and a note
is printed instead of raising.
draw_latex_inv(layout, tile_mn=None, filename="invlayout.tex", compile_pdf=True, color=white, cell_label=str, axis_label=index_label)
is the LaTeX/PDF analogue of draw_svg_inv, visualizing val -> crd. It
accepts the same layouts and tile, takes the same functors, and raises the
same errors. The filename and compile_pdf arguments behave as in
draw_latex.
>>> from pycute.util import draw_latex_inv
>>> L = Layout((4, 8), (8, 1))
>>> draw_latex_inv(right_inverse(L), shape(L))
Saved as invlayout.tex
Saved as invlayout.pdfdraw_latex_tv(layout, tile_mn=None, filename="tvlayout.tex", compile_pdf=True, color=thread_color_8x, cell_label=tv_label, axis_label=index_label)
is the LaTeX/PDF analogue of draw_svg_tv: draw_latex_inv for a
thread-value layout, with the functors taking (tid, vid). Each cell shows
(by default) T<tid> / V<vid> for the (tid, vid) landing there, colored
by the color functor (color(tid, vid) -> (r, g, b), defaulting to
thread_color_8x). The tile_mn, color, cell_label, axis_label,
filename, and compile_pdf arguments behave as in draw_svg_tv /
draw_latex, and the errors are those of draw_svg_tv.
>>> from pycute.util import draw_latex_tv
>>> SM80 = Layout(((4, 8), (2, 2)), ((32, 1), (16, 8))) # (tid, vid) -> offset
>>> draw_latex_tv(SM80, tile_mn=(16, 8))
Saved as tvlayout.tex
Saved as tvlayout.pdfEvery drawer takes a color functor that maps a cell key to an (r, g, b)
tuple (each component in [0, 255]). There are two key shapes, matching the
two families of drawers:
- offset functors —
color(idx) -> (r, g, b)— fordraw_svg/draw_latex, keyed on the cell's integer offset, and fordraw_svg_inv/draw_latex_inv, keyed on the integervaldrawn there. - thread-value functors —
color(tid, vid) -> (r, g, b)— fordraw_svg_tv/draw_latex_tv, keyed on the(tid, vid)pair.
pycute.util.draw_colors provides a small
catalog of ready-made functors:
| Functor | Signature | Colors by |
|---|---|---|
white (default for the other drawers) |
either (*args) |
constant white |
index_grey_8x |
color(idx) |
greyscale, idx % 8 |
bank_color_8x |
color(idx) |
shared-memory bank, idx % 8 (light spectrum) |
bank_color_16x |
color(idx) |
shared-memory bank, idx % 16 (light spectrum) |
bank_color_32x |
color(idx) |
shared-memory bank, idx % 32 (light spectrum) |
thread_color_8x (default for the TV drawers) |
color(tid, vid) |
thread, tid % 8 |
value_color_8x |
color(tid, vid) |
value index, vid % 8 |
warp_color_8x |
color(tid, vid) |
warp, (tid // 32) % 8 |
constant(rgb) |
factory → either | constant rgb |
The three bank_color_* palettes are evenly-spaced subsamples of one shared
32-color light spectrum, so they stay mutually consistent (pick the modulus
that matches your bank count or how many distinct colors read best).
>>> from pycute.util import (draw_svg, draw_svg_tv, bank_color_32x,
... value_color_8x, warp_color_8x, white, constant)
>>> # Spot shared-memory bank conflicts: cells sharing a bank share a color.
>>> draw_svg(Layout((8, 4), (4, 1)), color=bank_color_32x)
Saved as layout.svg
>>> # Color a thread-value layout by which warp owns each cell:
>>> tv = Layout((8, 8), (E(0), E(1)))
>>> draw_svg_tv(tv, color=warp_color_8x)
Saved as tvlayout.svg
>>> # ...or by value index instead of thread:
>>> draw_svg_tv(tv, color=value_color_8x)
Saved as tvlayout.svg
>>> # Flat coloring (e.g. for a clean figure), via white or a constant:
>>> draw_svg(Layout((4, 8), (1, 4)), color=white)
Saved as layout.svg
>>> draw_svg(Layout((4, 8), (1, 4)), color=constant((255, 225, 180)))
Saved as layout.svgAnything callable works: to roll your own, just supply a function (or lambda)
of the matching signature that returns an (r, g, b) tuple, e.g.
color=lambda idx: (255, 175, 175) if idx % 2 else (175, 255, 175).
Every drawer also takes two labelling functors, each returning the text to
write. A "\n" in the text starts a new line, so a cell can carry a label of
several lines, as the thread-value cells do by default, and an empty text
writes nothing.
- cell labels write each cell —
cell_label(val) -> strfordraw_svg/draw_latex, given the cell's offset or, for aTensor, its element, and fordraw_svg_inv/draw_latex_inv, given the indexvalthat lands there;cell_label(tid, vid) -> strfordraw_svg_tv/draw_latex_tv, given the(tid, vid)pair that lands there. An empty label leaves the cell blank, still colored. - axis labels write the margins —
axis_label(axis, i) -> strfor every drawer, given rowi(axis == 0, down the left edge) or columni(axis == 1, along the top). An empty label leaves that row or column unlabelled.
pycute.util.draw_labels provides the
defaults:
| Functor | Signature | Writes |
|---|---|---|
str (default for the other drawers) |
cell_label(val) |
the value |
tv_label (default for the TV drawers) |
cell_label(tid, vid) |
T<tid> over V<vid> |
index_label (default for every drawer) |
axis_label(axis, i) |
the index i |
>>> from pycute.util import draw_svg, draw_svg_tv
>>> # Rows and columns named by mode, m0..m3 and n0..n7:
>>> draw_svg(Layout((4, 8), (1, 4)), axis_label=lambda axis, i: f"{'mn'[axis]}{i}")
Saved as layout.svg
>>> # A thread-value layout as a register file: a row per thread, a value per cell.
>>> tv = Layout((4, 6), (E(0), E(1)))
>>> draw_svg_tv(tv, cell_label=lambda tid, vid: f"V{vid}",
... axis_label=lambda axis, i: f"T{i}" if axis == 0 else "")
Saved as tvlayout.svgA complete debugging session using the visualization toolkit:
>>> from pycute import *
>>> from pycute.util import *
>>> # Build a 4x8 identity tile (every cell is its own (m, n) coordinate):
>>> tile = Layout((4, 8), (E(0), E(1)))
>>> # Print as ASCII offsets:
>>> print_tensor(Layout((4, 8), (1, 4)))
(4, 8):(1, 4)
0 4 8 12 16 20 24 28
1 5 9 13 17 21 25 29
2 6 10 14 18 22 26 30
3 7 11 15 19 23 27 31
>>> # Save SVG visualizations:
>>> draw_svg(Layout((4, 8), (1, 4)), "data_layout.svg")
Saved as data_layout.svg
>>> draw_svg_tv(tile, filename="tile_coords.svg")
Saved as tile_coords.svg
>>> # ...or as LaTeX/PDF (same picture, vector output for papers):
>>> draw_latex(Layout((4, 8), (1, 4)), "data_layout.tex")
Saved as data_layout.tex
Saved as data_layout.pdf- Source:
pycute/util/print_tensor.py,pycute/util/print_table.py,pycute/util/draw_svg.py,pycute/util/draw_latex.py,pycute/util/draw_colors.py,pycute/util/draw_labels.py. - Tests:
test/test_draw.pyasserts on what the drawers write — labels, fills and sizes, and the inverse and thread-value views. The printers are not unit-tested, but you can driveprint_tensorfrom any layout-algebra test as a quick visual sanity check.
Copyright (c) 2023 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. SPDX-License-Identifier: Apache-2.0