A pure-Python reference implementation of CuTe — the hierarchical layout-and-tensor algebra at the heart of CUTLASS 3.x and the CuTe DSL. No GPU required.
Where C++ CuTe is a header-only template library tightly coupled to CUDA,
PyCuTe is plain Python you can import
from any script — making it the place to learn the algebra, prototype new transformations, and generate test vectors for the C++ and DSL implementations.
It implements the layout algebra from the CuTe Whitepaper — coalesce
,
composition
, complement
, logical_divide
, logical_product
,
right_inverse
, left_inverse
, nullspace
, recast
, layout_add
, and
greatest_common_domain
— over integer and coordinate (ArithTuple
/basis)
strides, plus limited support for F2
(XOR-swizzle) strides. A thin
Tensor
/Accessor
layer provides a reference data model.
Cris Cecka,
CuTe Layout Representation and Algebra,[arXiv:2603.02298]. The Whitepaper is the authoritative source for every definition and post-condition; PyCuTe defers to it throughout.
PyCuTe is installed in-place from a source checkout. Because many systems mark the system Python as externally managed (PEP 668), the recommended path is a virtual environment:
python3 -m venv .venv
source .venv/bin/activate
pip install -e . # core layout algebra only (no third-party deps)
pip install -e ".[viz]" # + visualization helpers (svgwrite, tabulate)
pip install -e ".[test]" # + everything needed to run the test suite
Python 3.10+ is the only hard requirement, and the core algebra has no
third-party dependencies. The optional extras add svgwrite
/tabulate
(viz
), sympy
(symbolic
), and pytest
(test
); the draw_latex
helpers
need only a LaTeX install (e.g. TeX Live's pdflatex
) for their PDF step. You
can also use the package straight from a repository checkout without installing
— import pycute
works as long as the repo is on PYTHONPATH
.
>>> from pycute import *
>>> A = Layout((3, 4), (4, 1)) # 3x4 row-major matrix
>>> A(2, 3) # call the layout on a coordinate
11
>>> A(11) # a 1-D coordinate works too
11
>>> size(A), rank(A)
(12, 2)
>>> coalesce(Layout((2, (1, 6)), (1, (6, 2))))
Layout(12, 1)
>>> composition(Layout(12), Layout((4, 3)))
Layout((4, 3), (1, 4))
>>> logical_divide(Layout(24), Layout(4, 2))
Layout((4, (2, 3)), (2, (1, 8)))
Build a Tensor
and read/write data:
>>> T = make_tensor(Layout((4, 4), (4, 1))) # 4x4 row-major
>>> T[1, 2] = 42.0
>>> T[1, 2]
42.0
Print or draw a layout (see Visualization):
>>> from pycute.util import print_tensor, draw_svg
>>> 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
>>> draw_svg(Layout((4, 8), (1, 4)))
Saved as layout.svg
The whole of CuTe fits in one sentence:
A
is a function from coordinates to offsets, defined by aLayout
(the coordinate(s) domain) and aShape
(how coordinates become offsets) of the same hierarchical profile.Stride
The stride is what turns a shape into a row-major, column-major, or arbitrarily nested map:
| Layout | Description |
|---|---|
Layout((4, 8), (8, 1)) |
|
| 4×8 row-major | |
Layout((4, 8), (1, 4)) |
|
| 4×8 column-major | |
Layout(((2, 4), 8), ((1, 16), 2)) |
|
| hierarchical (nested modes) |
A small algebra combines layouts to express tiling, partitioning,
vectorization, and layout analysis. Every operation is a pure function that
takes layouts and returns another Layout
:
— simplify to the fewest modes with the same map.coalesce(A)
— functional composition; indexcomposition(A, B)
A
throughB
.— the "missing" modes that fill outcomplement(A)
A
's codomain.— factorlogical_divide(A, T)
A
into tiles of shapeT
(tiling).— replicatelogical_product(A, B)
A
's pattern acrossB
(repetition).— invert and analyze maps.right_inverse
/left_inverse
/nullspace
A ** Tensor** is a
Layout
paired with an (e.g. a pointer): evaluating it at a coordinate evaluates the layout to an offset and dereferences the accessor at that offset. That short story is the whole of CuTe — everything else is a refinement or application of it. For the careful treatment of each piece, read the
Accessor
Start with docs/index.md. The documentation builds up from hierarchical tuples to layouts to the full algebra, with runnable examples drawn from the unit tests:
| File | Topic |
|---|---|
docs/00_quickstart.md |
docs/01_htuple.md
docs/02_shape_stride.md
Shape
, Stride
, and the integer-modules strides live indocs/03_layout.md
Layout
: construction, evaluation, coordinates, slicingdocs/04_layout_algebra.md
docs/05_tensor.md
Tensor
and Accessor
docs/06_swizzle.md
Swizzle
and F2
-stride layoutsdocs/07_visualization.md
print_tensor
, draw_svg
, draw_latex
, and color functorsdocs/08_api_reference.md
PyCuTe renders layouts as ASCII tables (print_tensor
, print_table
), colored
SVGs (draw_svg
, draw_svg_tv
), or TikZ/PDF (draw_latex
, draw_latex_tv
—
the analogue of cute::print_latex
). Every figure below is a plain PyCuTe
Layout
drawn with pycute.util
; regenerate them all with examples/readme_figures.py:
python -m examples.readme_figures # writes docs/images/*.svg (needs the viz extra)
A layout is a shape plus a stride. The same 8×8 shape with two different
strides gives a row-major or a column-major map; each cell is labeled with its
offset and colored by offset % 8
:
Layout((8,8),(8,1)) Row-major |
Layout((8,8),(1,8)) Column-major |
Layout(((4,2), (2,4)), ((1,32), (4,8))) Blocked |
Thread-value layouts. draw_svg_tv
shows how a warp's (thread, value)
pairs tile a matrix — here the C-accumulator of an SM80 16×8
MMA, colored by thread. This is precisely the partitioning the layout algebra produces:
Layout(((4,8),(2,2)), ((32,1),(16,8)))
(tid, vid) → (m, n) in a 16×8 tile
Swizzles are just F2 strides. Coloring a shared-memory tile by bank (
bank_color_8x
) makes conflicts visible. A row-major tile stores every column
in a single bank (vertical stripes → 8-way conflict); swapping the integer
strides for F2
(XOR) strides permutes each row so that every column spans all
8 banks — conflict-free — with no special-casing in the algebra:Layout((8,8),(F2(1),F2(9))) Swizzled Column-major |
Layout((8,8),(F2(9),F2(1))) Swizzled Row-major |
See docs/07_visualization.md for every drawer, the
(r, g, b)
color-functor catalog (index_grey_8x
, bank_color_32x
,
thread_color_8x
, …), and the LaTeX/PDF output.
pycute/
├── docs/ # documentation (start at docs/index.md); figures in docs/images/
├── examples/ # standalone scripts (einsum, TV-layout, README figures)
├── test/ # pytest unit tests (one test_*.py per operation)
└── pycute/ # the importable package
└── util/ # optional printing and visualization helpers
The suite uses pytest. Install the test
extra (see Installation) and run it from the repository root:
pytest # the whole suite, quiet
pytest --log-cli-level DEBUG # with live logging
pytest test/test_coalesce.py # a single module
pytest -k coalesce # tests matching a keyword
Cris Cecka.
CuTe Layout Representation and Algebra.arXiv:2603.02298 - Jack Carlisle, Jay Shah, Reuben Stern, Paul VanKoughnett.
Categorical Foundations for CuTe Layouts.arXiv:2601.05972 - Yang Shi, U. N. Niranjan, Animashree Anandkumar, Cris Cecka.
Tensor Contractions with Extended BLAS Kernels on CPU and GPU.HiPC 2016, pp. 193–202 - Bastian Hagedorn, Bin Fan, Hanfeng Chen, Cris Cecka, Michael Garland, Vinod Grover.
Graphene: An IR for Optimized Tensor Computations on GPUs.ASPLOS 2023, pp. 302–313 - NVIDIA CUTLASS / CuTe (C++)— the original C++ implementation. - NVIDIA CuTe DSL— the Python DSL that JIT-compiles CuTe kernels.
Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. SPDX-License-Identifier: Apache-2.0