Skip to content

Add an experimental n-dimensional Tensor (gb.Tensor, gb.einsum) - #651

Open
eriknw wants to merge 161 commits into
51-coverage-main-gapsfrom
50-tensor
Open

eriknw wants to merge 161 commits into
51-coverage-main-gapsfrom
50-tensor

Conversation

@eriknw

@eriknw eriknw commented Oct 27, 2026 •

Copy link
Copy Markdown
Member

An experimental n-dimensional sparse array, gb.Tensor, and gb.einsum, built on the
existing Matrix machinery. Both are marked experimental in their docstrings and docs, and the
API may change. They need the SuiteSparse backend and SuiteSparse:GraphBLAS 10 or later (the
Container API); elsewhere a Tensor raises NotImplementedError and the Tensor tests skip.
import graphblas imports neither the Tensor nor Numba; both load on first use.

Data model. A Tensor is one Matrix, stored by row, plus a layout: an ordered partition of
the logical axes into a row group and a column group, each raveled in C order into one
dimension of the Matrix. Hypersparse storage keeps the raveled dimensions cheap. Shape
operations (permute_dims, .T, expand_dims, squeeze, and reshape where the layout
allows) change only metadata, with copy-on-write between views. An operation that needs
another layout gets it from a planner, which prices reshape, transpose, permutation and
rebuild steps and takes the cheapest sequence.

Semantics. GraphBLAS missing-value semantics (no fill value), with numpy / Array API rules
for shapes, axes and broadcasting, and python-graphblas's lazy expressions, masks,
accumulators and <<. When a Matrix, Vector, Scalar or ndarray meets a Tensor in an infix
operator, Tensor rules apply. Matrix and Vector methods still refuse a Tensor operand, with a
hint.

What it covers

  • construction and conversion: from_coo, from_dense, from and to Matrix and Vector, and
    pydata/sparse and scipy n-d arrays in both directions (n-d scipy arrays need scipy 1.15);
  • element-wise operations with broadcasting, apply, select, comparisons;
  • reductions with monoids and aggregators;
  • contractions: tensordot, matmul / @ with batch axes, einsum with any semiring, TTV,
    TTM and MTTKRP. TTM and MTTKRP have selectable engines: GraphBLAS products, generated Numba
    kernels, and an array-UDT engine that uses the SuiteSparse JIT;
  • outer indexing, index assignment, take, take_along_axis and vindex;
  • reshape / ravel / flatten, relayout, unfold.

Changes outside the Tensor, each in its own commit (candidates to split into their own PRs):

  • Core hooks. Infix operators on Matrix, Vector and Scalar return NotImplemented for a
    Tensor operand, so Python tries the Tensor's reflected operator. Masks combine across kinds.
    Operator calls (binary.plus(A | T), unary.abs(T), agg.sum(T)) dispatch to the Tensor.
    A Tensor given to a Matrix or Vector method, as a mask of one, or to << raises with a hint.
    Core modules see the Tensor only through utils._tensor_types, which stays empty until
    graphblas.core.tensor is imported.
  • Matrix.isequal / isclose and Vector.isequal / isclose accept a Tensor.
  • The gb.Tensor and gb.einsum exports (lazy, as for Matrix, Vector and Scalar), and a
    tensor_max_expansion config key (default 1e9 entries): a broadcast in a union that would
    expand past it raises unless called with allow_expansion=True.
  • CI: the notebook job runs the Tensor notebooks only on SuiteSparse:GraphBLAS 10 or later.

Two changes that began in this PR now live in #652, below it: JIT C definitions for lifted ops on
a scalar and a UDT or on two UDT types (the Tensor's array-UDT engine uses them), and
ss.fix_jit_config finding a conda compiler on Windows.

Also here: a user-guide page (docs/user_guide/tensor.rst, with doctested examples) and the
API reference; ai-notes/, a new directory of curated notes for contributors and reviewers
(design, decisions, verification, limitations and roadmap, benchmarks, and a SuiteSparse
wishlist); asv benchmarks (benchmarks/tensor_ops.py); a portable comparison kit
(scripts/tensor_bench/: pydata/sparse, scipy, torch and Numba, plus optional SPLATT, TACO and
Finch builds); and two notebooks.

Where to start reviewing. The user-guide page shows the API. ai-notes/tensor/design.md
explains the data model, the planner and the engines, and has a map of the modules. The parts
that deserve the hardest look:

  • copy-on-write between O(1) views (_Owners, _Store in tensor.py);
  • the Container surgery in tensor_layout.py, which edits SuiteSparse's arrays in place;
  • arrays from the Numba kernels handed to SuiteSparse without a check (matrix_from_arrays),
    and, with suitesparse-graphblas before 10.5.1.1, handed over by utils._give_buffer's
    fallback, which follows the ss pack methods;
  • the planner's cost constants, fitted on one machine (an Apple M5 Pro).

Verification. The unit tests compare with dense numpy throughout (test_tensor.py,
test_tensor_io.py). Randomized oracles, each with a check that a planted error is caught, run
as slow tests (test_tensor_oracles.py, --runslow). ai-notes/tensor/verification.md has
the recipes for running without Numba and with the JIT off.

Local runs (macOS arm64, --runslow, cold JIT and kernel caches) found five problems that CI
would have hit, and CI's first runs found more: scipy before 1.15 (no n-d arrays) and before
1.13 (no coords), pydata/sparse 0.14, a test that checked the wrong backend, a CI draw with no
Numba, Windows' 32-bit default int in doctest output, and a dataset whose server rebuilt its
zip. The branch's last commits, from "Tensor tests: skip the Numba tests on any ImportError" on,
fix them, and CI is green. With those commits, these pass locally: the full suite
with suitesparse-graphblas 10.5.1.1 and with 10.3.1.0, each also with CI's C-scalar second
pass; the full suite without Numba, on both backends; the Tensor modules with the JIT off;
Python 3.11 with numpy 1.24 and no C compiler; scipy 1.13 with pydata/sparse 0.14; the Tensor modules in two shuffled orders; the
user guide's doctests; the Quickstart notebook; and pre-commit run --all-files. Three tests
from the base branch fail in some of those environments, as they do without this PR (one with
the JIT forced off, one with no C compiler, one without installed package metadata). The
oracles take about 90 s with cold caches here, and the Tensor modules about 4.5 minutes of an
11-minute run. Not run locally: Linux, Windows, x86, and SuiteSparse:GraphBLAS before 10 (CI covers them). On
CI's slowest runner (macos-15-intel) the Tensor modules take about 20-35 minutes, mostly Numba and
JIT compiles.

Review and coverage since the first CI runs. An adversarial review looked for what CI and the
oracles miss and found 18 issues: crashes (take_along_axis over a broadcast size-1 axis read out
of bounds; Container work on a Matrix with pending work; Numba kernels with 64-bit row pointers),
wrong results (C << T.T when C shares T's storage; reductions of an empty region of an iso
Tensor, from a long-standing SuiteSparse bug that Matrix also has; einsum and MTTKRP regrouping
with non-distributive or non-commutative semirings; broadcast integer pow), shapes past 2^60 and
2^63 positions (broadcasting, TTV, extraction and gather costs), and smaller API gaps (-T,
abs(T), ~T, string select, GCXS with repeated columns, a C-name collision with a user UDT).
Each fix has a regression test that fails before it, and the oracles now draw the cases they
missed. A second review, by an agent that wrote none of the fixes, checked them (all correct; its
notes are fixed in the last commits). A coverage pass then brought the Tensor's source lines to
about 99.8% covered locally (over the suitesparse, no-Numba, JIT-off and vanilla runs), with tests
first, 17 unreachable branches removed, and pragmas only on the Numba kernels' definitions and two
safety guards. A performance round followed (each change reviewed the same way): extraction is priced
by the rows it visits, a row becomes a Vector without an extract, a column range of sparse rows
is filtered without an extract, one planner search prices many targets (a new 5-d MTTKRP shape
plans in 12 ms, not 163), and unsigned sums use a precompiled kernel. Uber's slices went from
0.52x to 1.19x of the fastest of pydata/sparse, scipy and torch at 1 thread (1.68x at 8). Two
choices from it are recorded in ai-notes/tensor/decisions.md (24, 25): wait() keeps
building SuiteSparse's row hash, as Matrix.wait does (wait("complete") defers it), and
there are no per-mode copies for TTM; an explicit relayout() before a loop does that job.

Not tested yet, and known gaps (ai-notes/tensor/limitations-and-roadmap.md): before this
PR's CI, the Tensor had run only on macOS (arm64) with SuiteSparse:GraphBLAS 10.5.1 and 10.3.1,
so Linux, Windows and other SuiteSparse versions are new to it; the planner's constants are unmeasured on x86;
TTM over axes off the stored order is 2-7x behind TACO, which keeps a copy per mode. The
comparisons with pydata/sparse, scipy, torch, SPLATT, TACO and Finch are in
ai-notes/tensor/benchmarks.md.

History. The commits are the development history, kept as written. This repository's squash
merge puts every commit message into the merge commit by default, so the merge message wants
replacing with a summary.

🤖 Generated with Claude Code

@eriknw
eriknw added this pull request to stack #650 October 27, 2026 04:30
@codecov

codecov Bot commented Oct 27, 2026 •

Copy link
Copy Markdown

Codecov Report

❌ Patch coverage is 97.90295% with 121 lines in your changes missing coverage. Please review.
✅ Project coverage is 98.75%. Comparing base (71999d8) to head (6f0ebd6).

Additional details and impacted files
@@                    Coverage Diff                    @@
##           51-coverage-main-gaps     #651      +/-   ##
=========================================================
- Coverage                  99.30%   98.75%   -0.56%     
=========================================================
  Files                        109      125      +16     
  Lines                      33019    43389   +10370     
  Branches                    3734     5130    +1396     
=========================================================
+ Hits                       32791    42849   +10058     
- Misses                       152      401     +249     
- Partials                      76      139      +63     
🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

eriknw added 15 commits October 28, 2026 20:45
The core of the Tensor (data model, relayout planner, metadata-only
shape ops with copy-on-write, recipe expressions with masks and accum,
same-shape element-wise ops and reductions). New modules only:

- graphblas/core/tensor_layout.py: Layout, raveling, Container helpers
  (by-col to by-row flip, nonempty rows, Vector to row), a Dijkstra
  relayout planner over (seq, split) states with RS / T / PERM_R / REBUILD
  edges, plan and selector caches.
- graphblas/core/tensor.py: Tensor (BaseType subclass), TensorExpression
  (always a recipe), from_coo/from_matrix/from_vector, to_coo/to_matrix,
  permute_dims/T/mT/expand_dims/squeeze, ewise_add/mult/union, apply,
  reduce/reduce_scalar/sum/prod/min/max/any/all, infix + - * /.
- graphblas/tests/test_tensor.py: numpy model with a present mask.
Pass only the user's descriptor options to internal updates (building a
descriptor per call cost about 20 us), and look up selector semirings
only when a selector is used. Move a reduction's result to the default
split only when its column group is empty: a result held as one row is
as fast as a square Matrix for later ewise, apply and selector ops, and
the reshape cost 0.7 ms of 2.0 ms for a sum over the last axis at 1e7
nnz.

Add tests for a UDT relayout and for size-0 axes.
Relayout each operand so its contracted axes are exactly one group, in
the same order on both sides, choosing the contracted order and the
free-axis orders by the planner's estimated cost, then one mxm with the
semiring. The result layout is (free axes of self | free axes of other)
in whatever order the planner picked. Tests compare with numpy.tensordot
on the numpy model, and min_plus by brute force.
…in place

- The hypersparse branch (REBUILD with dup_op) replaced even a rowwise
  reduce, which is cheap at any density: a sum over the last axis of the
  (2^19)^3 case took 6.8 ms instead of 0.7. Check density only before a
  selector or columnwise step.
- A wide columnwise reduce goes to a selector only when the selector is
  small next to the data; on FROSTT Uber a selector of 5e6 positions for
  3.3e6 entries took 32.8 ms vs 4.9 for a one-thread columnwise reduce.
- REBUILD after a partial reduction takes the kept axes from the current
  layout.
- Relayout reshapes its own intermediates in place (7.1 ms vs 9.6 for 1e7
  entries, less peak memory), except when the new dimensions pass 2^31:
  SuiteSparse 10.5.1's in-place reshape then gets column indices wrong
  (they look truncated to 32 bits), which test_relayout_hypersparse caught.
Up to 32 cached permutation and selector matrices of up to 2^24 entries
each could hold several GB. Weigh each by its entries and keep at most
2^25 in all (about 400 MB).
Core hooks (the only edits outside the Tensor modules): infix operators,
_matmul_infix_expr and call_op on Matrix, Vector and Scalar return
NotImplemented when a Tensor is involved, so Tensor rules apply through
the Tensor's reflected operator; in-place Matrix/Vector operators and
assignment with a Tensor raise TypeError; a Tensor mask on a Matrix or
Vector output raises; op(T, x), unary(T), agg(T) and select ops dispatch
to Tensor methods; _expect_type_message hints at Tensor methods. Tensor
types register in utils._tensor_types, so core modules never import the
Tensor module.

Tensor: Matrix (2-d), Vector (1-d) and ndarray (full) operands; numpy
broadcasting for ewise_mult and comparisons by diagonal scaling (exact
group, part of a group through a selector, or a planned relayout) and
for ewise_add/ewise_union by replication behind an O(1) size guard
(config tensor_max_expansion, default 1e9; allow_expansion=True opts
in); comparison infix; |, & and @ infix expressions for op(X | Y);
value select; from_dense/to_dense.
…onal embedding

matmul follows numpy: the last two axes multiply, leading axes broadcast,
and a 1-d operand is promoted and its axis removed. Batch axes that only
one side has are free axes of that side (plain tensordot); batch axes both
have lead both row groups, and the left operand's columns are shifted per
batch in the Container's own index array (keeping 32-bit indices, which
ss.unpack would widen to 64 bits), so one mxm does all batches. Only the
left side is embedded, so the product keeps the right side's columns and
needs no unembedding; it beat a per-slice loop at every batch count at 8
and 1 threads, so there is no per-slice path. semiring(X @ Y) works
through an infix expression, which now checks shapes when built, as
Matrix does.
einsum(subscripts, *operands, op=semiring) lowers pairwise: indices only
one operand has (and nothing later needs) are reduced first; shared
indices are batch (kept) or contracted; no contraction is an element-wise
product with broadcasting; contraction goes to tensordot, or with batch
indices to the block-diagonal product. Three or more operands follow
opt_einsum.contract_path when it is installed, else a greedy order. The
MTTKRP pattern goes to Tensor.mttkrp (one TTM, then one product per
remaining axis with an (I R x R) embedding of its factor), since
opt_einsum would form the Khatri-Rao product first.

tensordot with a 1-d operand over an axis inside a group of the stored
layout scales by a diagonal and reduces with a selector; over the
column group it is one mxv. unfold(mode, order='C'|'F') gives the
numpy/tensorly or Kolda-Bader matricization.
T[key] extracts with outer (orthogonal) semantics, as GraphBLAS extract:
ints drop their axis; slices, lists and int arrays keep it; ... fills in.
Each group of the stored layout becomes one extract index: all, a
GraphBLAS range when the selection is a contiguous run of raveled
positions (a leading int or range with the rest of the group whole, so
slabs of hypersparse groups stay cheap), else a listed product; past
2^26 listed positions it filters the coordinates instead.

from_matrices stacks equal Matrices along a new axis with one
ss.concat (the axis only picks the layout labels); to_matrices splits
with ss.split after one relayout.
…axis order

TTV over an axis inside a group of the stored layout now folds the vector
into the selector (S[kept, g] = w[axis coordinate of g]) and runs one
product over the data, instead of a diagonal scaling then a selector
reduce: 1000^3 / 1e7, mode j 5.7 ms vs 9.9, mode i 10.3 vs 14.2 at 8
threads. MTTKRP tries each order of the other axes and takes the one
whose layout is cheapest to reach.
…out Numba

binary.floordiv is a UDF without Numba, and resolving it at import made
graphblas.core.tensor fail to import in the no-Numba CI configuration.
SuiteSparse 10.5.1's in-place reshape cannot widen 32-bit indices; an
intermediate past 2^31 columns (J * R) now takes the copying reshape.
…a strategy

ewise_mult with broadcasting took an operand whose broadcast factor is 1
as the one covering the result. A lower-rank operand whose missing
leading axes are all size 1 also has factor 1, so (B * A) with B of
shape (1, 4) and A of shape (1, 1, 4) used B's 2-d layout against the
3-d result; a REBUILD then left one axis without coordinates and failed
(found by an independent check against a numpy model). Both operands
now get leading size-1 axes up to the result's rank first (metadata
only), and relayout rejects a layout whose rank does not match the
shape.

Adds a sweep over every shape pair with sizes {1, 2} up to 3-d (lower
rank and 0-d included), both operand orders, several layouts, for
intersection, union and union-with-defaults; it fails without the fix.
…build P from the Container

PERM_R was priced per entry plus min(nrows, nnz) rows, but rows that are
mostly empty cost a hash lookup each in P @ M (about 60 ns at 8 threads),
and P was built per call with np.unique over every entry's row index.
On skew (1e5, 1e3, 10) at 1e7, permute (1, 0, 2) was predicted at 26 ms
and took 3.5 s. Now the planner prices PERM_R by the expected nonempty
rows of the state (uniform occupancy), with the hypersparse per-row cost
when rows are mostly empty, plus building P (numpy per row when P is
too big to cache, per nonempty row when P comes from them). P is built
from the nonempty rows read from the Container (h, or diff(p)) in
O(rows) instead of np.unique (3.0 s -> 0.14 s for 9.5e6 rows).

skew perm_102: 3,507 -> 34 ms at 8 threads (REBUILD 142, torch 801);
NELL-2 perm_021 9,939 -> 888 ms (torch 2,844), perm_102 1,385 -> 714.
… h; TTV by vxm

The first benchmark run found NELL-2's default (0|1,2) layout behind by
100x on sums, TTV and broadcasts along axes inside its 2.6e8-position
column group: a shape-only selector costs per position, not per entry,
and was allowed up to 4 x nnz positions.

- The default split now follows the data: from_coo builds with the most
  axes in the rows and moves the split left (one reshape each) while the
  rows hold fewer than 1.5 entries per nonempty row, counted through the
  Container. A row group indexes its nonempty rows (h), a column group
  does not, so sparse axes belong in the rows. NELL-2 becomes (0,1|2).
- Selectors over a hypersparse row group (reductions, weighted TTV,
  broadcast expansion) are built from its nonempty rows, sized by the
  data. Column-group selectors are priced per position and lose to a
  REBUILD with dup_op (or a split move) when that is cheaper.
- TTV picks the cheapest of mxv (axis is the column group), vxm (the row
  group: one product, no relayout; relayout + mxm took 56 ms vs 36), a
  weighted selector, a split move, or the general relayout and mxm.
- Broadcast multiply picks between an exact group, an expansion over a
  group (from h for hypersparse rows), a split move, or a relayout.

NELL-2 at 8 threads, first benchmark run -> now: sum over j 1572 ->
70 ms, over k 1653 -> 2.8, TTV j/k 4097/4014 -> 76/4.9, X * v[j]/v[k]
1416/1504 -> 41/20; sum over i 699 -> 82; X * v[i] 15 -> 41 (torch 84).
A (3000, 2000, 10000) tensor of 7.7e6 entries in its default layout:
sum(axis=2) 50 -> 4.6 ms, TTV k 237 -> 4.7, X * v[k] 19 -> 2.8; in
(0|1,2), TTV over the row group 56 -> 36 ms (vxm 35).
eriknw added 28 commits October 28, 2026 20:45
…elector

The full weight vector added to test_hypersparse_row_group_paths ran by
another option; force the row selector (where it is offered), whose
weights then keep every nonempty row.
… side at the entries

Past 2^60 positions in total, a broadcast ewise_mult whose op has no flipped form of its kind
(pow: rpow is a Numba UDF) and whose small side is sparse found no option in _bcast_mult: no
gather (the small side is not full), no diagonal side (row_op needs the flip, col_op's
groups would pass 2^60), no split moves and no planned relayout (past 2^60 neither fits).
It then replicated the small side over the broadcast axes, as for a union, and _expand
could not lay out a 2^80-row operand: as_layout's "more than one Matrix dimension" error,
while the same product with the operands swapped worked.

That fallback now looks the small side's values up at the big side's entries: their
coordinates along the small side's axes, from the stored rows and columns, through the
gather's lookup (tensor_index.lookup, factored out of gather), as a Matrix of the big side's
pattern, then one ewise_mult in the operands' order. It costs by the entries, so it is
right for an intersection at any size.

A union replicates by its nature. _expand now puts the broadcast axes after the column
group when they do not fit before the row group, and raises a clear ValueError when they
fit in neither, or when the copies would pass 2^60 entries (with allow_expansion=True;
without it the expansion limit refuses first, as before). design.md says so.

A new test (with the Numba and the GraphBLAS lookups) runs pow broadcast both ways on a
2^120-position Tensor of 3 entries against sparse small sides along one axis (the lookup)
and two axes (a relayout fits), checks the census choice, the union's refusals, a small-factor
union whose broadcast axis goes to the columns, and the error when neither group can take it.
The Numba TTM's output Matrix has a row for every combination of the free axes (the stored
rows, then the leading column axes), and _ttm_numba did not check their count. With Numba,
2^16 entries or more, and free axes past 2^60 positions (X of shape (2^31, 2^31, 4) stored
as (0|1, 2), TTM over axis 2), tensor_numba.ttm built a 2^62-row output, which GraphBLAS
refused with a bare InvalidValue(''). Without Numba, or with fewer entries, the products
raised "the contracted or free axes do not fit in a Matrix dimension (2**60)"; the array-UDT
engine declines such a case (it needs the free axes as one group).

_ttm_numba now declines when the free axes pass MAX_GROUP positions, before planning a
relayout, so the products give their clear error.

A new test, with the kernel's entry minimum off, checks the ValueError and that the kernel
was not called for (2^31, 2^31, 4), and that for (2^31, 2^29, 4) (free axes of exactly 2^60)
the kernel still runs and gives the model's values.
Past 2^60 positions in total, _ttv has no split moves and no general relayout (the other
axes as one group do not fit), so when the axis lies inside a column group of several axes
its one option was "colsel": a shape selector over the whole group, whose time and memory
grow with the group's positions however few the entries. For a (2^31, 2^31, 4) Tensor
stored as (0|1, 2), two entries took over 2 minutes and about 10 GB (the selector over 2^33
positions). Below 2^60 that selector loses to a REBUILD or a split move (6bd3586's per-
position pricing); past it nothing else was offered. With an op that has no flipped form
(plus_pow, w first) there was no option at all, and the products refused ("do not fit").

_ttv now also offers "entries" past 2^60 (_ttv_by_entries): w's value at each entry's
coordinate along the axis, through the gather's lookup, as a Matrix of the Tensor's pattern;
one ewise_mult in the operands' order (no flip needed); then the axis reduced by the monoid
with _reduce_matrix, which takes a REBUILD where a selector would cost more. It costs by the
entries: the case above takes a millisecond and 110 kB. It is priced as a lookup plus a
REBUILD, and competes only past 2^60, so nothing changes below. With "entries" always there
past 2^60 and the general option always there below, the "no option" branch could not run
any more; it is gone, and the coverage test that reached it (plus_pow, w first) now checks
the result through the entries. design.md's TTV section says so.

A new test, with column groups of 2^24 and 2^33 positions, runs TTV with plus_times and
plus_minus, the vector on either side, the axis in the column group (entries) and in the
row group (the row selector over the nonempty rows), against a model, under a tracemalloc
peak of 16 MB (2^24 on the old code: 1.2 GB; 2^33 was not run there).
…hint

Since TensorIndexExpr is registered as giving a Tensor (for -T[i], e3dc693, finding 14),
_as_tensor takes an extraction for Tensor-like; but _compute turned only an exact
TensorExpression into its value, so T + T[0], T.ewise_add(T[1]), tensordot and einsum with
T[i] (and, with autocompute, T[0] * T[1] and binary.plus(T[0] | T[1])) passed the
uncomputed extraction on as if it were a Tensor and raised AttributeError ("no attribute
'_M'"). Before that registration they raised a TypeError. No value was wrong.

_compute now tests isinstance(x, TensorExpression), which covers TensorIndexExpr: with
autocompute off the TypeError says to call .new() (the hint), with it on the extraction is
computed.

A new test runs seven such operations with autocompute off (the TypeError and its hint) and
on (numpy's results).
Tensor._update ran an exact TensorExpression and refused everything else that is not a
Tensor or a scalar, so C << T[1] and C(mask, accum) << T[1] raised "Assignment value must be
a Tensor, TensorExpression or scalar, not TensorIndexExpr", although an extraction is a
TensorExpression (a subclass) and Matrix's v << A[1, :] works. It has been so since the
Tensor came in.

_update now tests isinstance(expr, TensorExpression), so an extraction runs its recipe on
the target, with the mask and the accumulator, like any other expression.

A new test stores an int, a slice and an outer-indexing extraction into new Tensors, and
into Tensors with old values under a structural mask and a plus accumulator, against numpy.
einsum's allow-list (5be07de) regroups min_plus and max_plus on integers too: it sums an
index out first and contracts in steps. Integer plus wraps, so when a sum overflows the
regrouped result can differ from the term-by-term one: min_plus on INT8 "ij,jk->i" with
a = [[100]] and b = [[50, -100]] has terms that wrap to [-106, 0] (minimum -106), and the
Tensor gives 0, having summed b over k first. max_times was left off the list, with integer
overflow among the reasons given; the review found the two cases inconsistent.

They are kept on the fast path, and the docs now say what it gives: exact up to
floating-point rounding and integer overflow. A min_plus sum that overflows is no
meaningful distance either way (the wrapped -106 is not one), min_plus is a common semiring
whose term-by-term path forms every term (much more memory), and plus_times stays exact
when it wraps (wrapping arithmetic is a ring). max_times and min_times stay off for their
other reason: on signed values they do not distribute even without overflow. The allow-
list's comment, einsum's docstring and the user guide say so. Documentation only; no
behavior changes, so there is no new test.
Two comments from the coverage pass, both about code that is right:

- _ttm_numba's "(never empty ...)" gave the wrong reason. With 5 or more free axes only
  their natural order is enumerated, so the stored grouping may not be a target. What
  keeps the targets non-empty is the free-axes check above (d863d94): the free axes fit
  one group and the axis another, so (free axes | axis) is always a target. The comment
  says so now.
- _lookup_numba's by-column return was marked "pragma: no cover (safety: a Tensor's Matrix
  is held by row)", but it is reachable: a Matrix the Tensor shares (from_matrix with
  copy=False, or T.ss.matrix) that the user flips to by-column. The guard is kept, the
  pragma is gone, and the comment says when it runs.

A new test makes that state (sparse and bitmap, flipped with format "by_col"), gathers
through vindex, and checks that the Numba lookup declined and the GraphBLAS one gave the
right values. It passes on dff2457 as well: the guard was right, only its note was not.
SuiteSparse's extract finds each listed row directly (or by a hash lookup
when the Matrix is hypersparse), but the cost model charged it every row of
the Matrix, so slabs inside a sparse row group filtered the coordinates
instead. A column list under 1/32 of the columns is inverted into a hashed
matrix, and each entry of a visited row then pays a lookup; that is priced
too.

A/B against b45a6c48 over the kit's slice group on 8 data sets, interleaved,
2 reps, medians of 7 (load 6.1-6.7): 10 of 104 rows change strategy, each to
the fastest forced option (ms, 1 / 8 threads):

  Uber slice1   24.5 -> 2.28    14.0 -> 1.34
  Uber slice2   23.5 -> 0.125   13.2 -> 0.149
  Uber range1   36.2 -> 4.20    26.7 -> 2.70
  Uber range2   33.4 -> 4.59    23.9 -> 2.71
  u6 slice2     14.5 -> 4.62    (extract kept)
  Zipf slice2  103.8 -> 40.8    (extract kept)

The other rows are within noise (rechecked at load 4.2); fingerprints match.
Broadcasts and TTV made the small side's (1 x n) row a Vector by
M[0, :].new(), an extract that runs SuiteSparse's generic subref without the
JIT (SuiteSparse has no factory one), and Tensor.to_vector did the same.
tensor_layout.steal_row_as_vector, the inverse of steal_vector_as_row, takes
the row over in O(1). Where the relayout returned the operand's own Matrix,
it is copied first (dup), so the operand is left as it was.

A/B against b45a6c48 over the kit's bcast and ttv groups on 8 data sets,
interleaved, 2 reps, medians of 7 (load 4.5-7.7; ms, before -> after):

                            JIT on          JIT off
  NELL-2 mul_mat12, 1t   1124 -> 967     1509 -> 964
  NELL-2 mul_mat12, 8t    328 -> 276      364 -> 271
  u6 mul_mat12, 1t       4.29 -> 3.63    5.77 -> 3.56
  hyper mul_vec0, 8t     2.03 -> 1.62    2.12 -> 1.65

Geometric mean over the 112 rows: 0.98 (on), 0.975 (off); no row slower
beyond noise (rechecked); fingerprints match.
X[..., k] and X[..., k0:k1] along the column group of a sparse (not
hypersparse) stored Matrix ran SuiteSparse's extract, which visits every
row, empty or not: Uber's (0, 1, 2 | 3) has 5e6 rows for 3.3e6 entries. One
numpy pass over the column indices in place, with the kept entries' rows by a
binary search of the row pointers, is priced against it (per entry, and per
entry kept) and taken where cheaper. It also avoids the generic subref
without the JIT. Hypersparse rows are left to the extract, which visits only
the nonempty ones (32 entries or more in the default layouts).

A/B against b45a6c48 over the kit's slice group on 8 data sets, interleaved,
2 reps, medians of 7, with the extract pricing of the previous commit (load
4.6-5.2; ms, JIT on / off):

                     1 thread                    8 threads
  Uber slice3    31.8 / 32.8 -> 1.20 / 1.18   7.34 / 7.32 -> 1.21 / 1.20
  Uber range3    31.1 / 31.9 -> 1.34 / 1.35   11.9 / 12.8 -> 1.41 / 1.53
  u7 slice2      19.7 / 20.3 -> 3.53 / 3.55   (extract kept)
  hyper slice1   6.56 / 6.81 -> 0.48 / 0.43   2.05 / 2.17 -> 0.35 / 0.40
  hyper range1   7.40 / 7.61 -> 4.23 / 4.01   (extract kept)

Each changed row is within 10% of the fastest forced option; the others are
within noise (small slices rechecked with 41 calls); fingerprints match.
Choosing an MTTKRP's axis order prices a relayout to each candidate order
(24 in 5-d), and each was its own planner search over the states. Searches
from one source settle the same nodes in the same order up to their target,
so one search continued past all the targets gives each the cost and path its
own search finds; tensor_layout.plan_each does that and caches each result
under the key plan would use. A test compares the two on 3-d to 5-d shapes,
sizes and thread counts.

The first MTTKRP on a new shape (tiny tensors, plan cache cleared, each mode,
median over 6 shapes, 3 interleaved runs per tree, load about 5):

  5-d   163 -> 12 ms
  4-d   4.4 -> 2.2 ms
  3-d   0.38 -> 0.38 ms

The chosen orders are the same, and every census decision of the kit's
mttkrp, ttm, ewise and contract cases on u7, Uber, NIPS and NELL-2 at 8 and
1 threads (144 rows) is unchanged.
A row selector (S @ M) holds True and reduces with (monoid, second).
SuiteSparse compiles PLUS_SECOND_UINT32 and _UINT64 out of its factory set,
so sums of unsigned data ran its generic kernel without the JIT. A selector
of 1s of the data's type with plus_times is a factory kernel and gives the
same sums. Column selectors (M @ S.T) keep BOOL: with the JIT on, plus_first
does not read the selector and was 2x faster there than plus_times (0.87
against 1.73 ms for 1e6 entries at 1 thread).

A/B against b45a6c48 on u7 cast to UINT32, UINT64, INT64 and FP64, every sum
over one or two axes, interleaved, 2 reps, medians of 7 (load 3.3-3.8; ms,
JIT off, 1 / 8 threads):

  UINT32 sum(0)   63.2 -> 25.9    12.3 -> 6.26
  UINT32 sum(1)   55.7 -> 20.3    8.93 -> 3.88
  UINT64 sum(0)   64.3 -> 28.4    13.0 -> 6.98
  UINT64 sum(1)   53.0 -> 21.2    8.64 -> 4.11

INT64 and FP64, and every row with the JIT on, are within noise (two-axis
sums at 8 threads rechecked with 41 calls); results match.
… of all five

limitations-and-roadmap.md gets the breakdown of TTM along a non-stored axis
(relayout, kernel, and the hyper hash SuiteSparse builds for the output in
the wait() that ends each timing: 39-55% of NELL-2 modes 0 and 1 at 1
thread), the in-place kernel prototype that did not beat the relayout, and
the open question of whether Tensor.wait should build that hash.
benchmarks.md points there, and records the A/B of the five previous
commits together against b45a6c48: 762 kit rows, geometric mean 0.949, no
row slower beyond noise once rechecked.
steal_row_as_vector flipped its (1 x n) Matrix as if it were held by row.
A shared Matrix the user flipped to by-column (a state the gathers support)
then came out as an (n x 1) Matrix held by row: no valid Vector (size 1,
GxB_Vector_fprint -104) from Tensor.to_vector, and the TTV and broadcast
rows went through the same invalid object. The Matrix is now converted to
by-row first; it is always the caller's own (a dup or a relayout's new
Matrix), so nothing shared changes. Found by Review 4.
… control

- The column range filter reads the stored arrays where they are and copies
  the entries it keeps; "in place" read as if it changed them. Comments,
  docstring and notes now say so (the commit that added it keeps its subject).
- plan_each's argument and design.md: searches from one source settle the
  same nodes up to each target only while the target is within its own
  REBUILD cost; past it, both give the REBUILD. Said so.
- steal_row_as_vector made a hypersparse one-row Matrix sparse (a reshape's
  output keeps a hyper switch of 1.0), and the Vector kept that "sparse"
  control, so it could never become bitmap or full. The Matrix's own control
  goes back on the Vector, with "sparse" for its hypersparse part.
Two decisions, recorded with measurements (no behavior change):

- Tensor.wait() keeps materializing as Matrix.wait() does, including the
  hash SuiteSparse builds of a hypersparse Matrix's rows. Its new docstring
  says what that costs and what wait("complete") does instead: nothing, so
  the next operation pays it. On NELL-2's TTM along modes 0 and 1 (R = 16)
  the hash takes 0.68-0.97 s at 1 thread and 86-116 ms at 8, and is all the
  wait does there; after wait("complete") a row extract that took 1-3 ms took
  about the hash, and to_coo and a sum about the hash longer.
- No per-mode copies for now (memory, and coherence with copy-on-write and
  in-place updates); relayout() before a loop gives one. With the relayout
  done outside the timing, as TACO's copy is, those TTMs take 445 and 548 ms
  at 1 thread without the hash (TACO: 263 and 262), 120 and 158 at 8.
  To revisit for many-mode sweeps such as HOOI and Tucker.

decisions.md gets both (24, 25); benchmarks.md the numbers next to TACO's
and Finch's; limitations-and-roadmap.md the TTM entry and roadmap item 4.
- The hash threshold is the hyperlist's length against SuiteSparse's
  GxB_HYPER_HASH setting (1024 by default), not a count of nonempty rows:
  the docstring and decisions.md say so.
- After wait("complete"), most operations (not all) finish their operand
  first and so pay the hash next: said so in the docstring and the notes.
- Decision 24 grounds "materialize" in SuiteSparse's documented
  GrB_MATERIALIZE semantics (all work finished; usable as an input by
  several user threads at once), not in Matrix.wait's docstring.
- Tensor.wait returns the Tensor, as Matrix.wait returns the Matrix
  (it returned None); a test checks both modes.
Matrix's update turns an infix expression into its expression with _to_expr(): C << (A | B)
stores the union with lor and C << (A & B) the intersection with land, on BOOL operands only
(other dtypes raise "Bad dtypes for `x | y`", since they need an operator), and C << (A @ B)
the default semiring's product. Tensor._update took none of them: C << (X | Y) on BOOL
Tensors raised "Assignment value must be a Tensor, TensorExpression or scalar, not
TensorEwiseAddExpr". And Tensors had no ^ at all (Python's "unsupported operand"), where
Matrix's A ^ B is the union with lxor, for BOOL only.

Tensor._update now takes the Tensor infix expressions through a new _assigned(): | and &
give their lor / land expression on BOOL and refuse other dtypes with a TypeError like
Matrix's (naming the operator to give instead); @ gives its default product. X ^ Y, its
reflected and in-place forms are the union with lxor, refused for a result that is not
BOOL with Matrix's own message; a Matrix on the other side defers to the Tensor, as for the
other operators. The (X | Y).new() form keeps its default operator, as before. The user
guide's infix sentence says so; the core-hooks test that expected ^ to be unsupported now
expects Matrix's BOOL-only TypeError.

A new test compares C << with |, & and ^ on BOOL Tensors with Matrix, plain and under a
structural mask with a lor accumulator, checks the 3-d results against numpy, checks the
refusals on FP64 for | and & (Tensor and Matrix alike) and for ^, compares C << (X @ Y) on
FP64 with Matrix, and runs ^= and ^ with a Matrix on either side.
…ay unsupported

Round 3 (fbd1f27) let a broadcast union put the broadcast axes after the column group when
they do not fit before the row group, and raise a clear ValueError when they fit in
neither. One case is left: past 2^60 positions, axes that would fit only split across
both groups. Supporting it needs a kron on both sides and a layout choice in _expand, for
a rare case, so the decision (made on Erik's instruction to decide and document) is not
to, for now. decisions.md records it as 26, with why and when to revisit (a workload that
meets the error). limitations-and-roadmap.md lists it among the missing features.
…oadmap idea

The perf round found that the Tensor's operations finish their operands with wait()
(materialize) before Container work, which also builds SuiteSparse's hyper hash of a
hypersparse Matrix's rows, even where the operation would not use it (SuiteSparse's own
extract of a short row list does not build it). On NELL-2's TTM output that hash is
0.68-0.97 s at 1 thread and 86-116 ms at 8 (about 45 ns a row), and after
wait("complete") a 1-3 ms row extract took 0.68-0.96 s because its own wait built it.
Internal waits cannot use "complete", which does no work at all, since Container work needs
pending work finished (stored_format). limitations-and-roadmap.md now lists this as a known
performance gap with those numbers and the idea of a per-Matrix GxB_HYPER_HASH setting
around the wait (with the caveat that GB_wait frees an existing hash while it is off), and
adds it to the performance item of the roadmap. No code change.
On non-BOOL Tensors, (X | Y).new() and (X & Y).new() computed the default operator (plus,
times), where Matrix raises "Bad dtypes for `x | y`!"; and the other ways of using the
expression on its own fared differently: C << (X | Y) refused (de6bd07), while passing it
to a method, using it as an operand or calling to_dense raised AttributeError ("no
attribute '_shape'"), BOOL included. The coordinator decided (on Erik's instruction) to
align with Matrix: the user guide says infix operators work "as for Matrix"; with numpy,
| on integers is bitwise or and on floats raises, so computing plus surprises; + and *
already give plus and times; and nothing recorded a reason for the old default.

- _TensorEwiseInfixExpr._to_expr gives lor / land on BOOL and raises Matrix's own "Bad
  dtypes" TypeError (its text, from infix._ewise_add_to_expr) for other dtypes. C << takes
  it directly (no separate _assigned).
- _compute turns an infix expression into its expression first, so as an operand (ewise,
  infix, einsum, region values through _region_value) it raises for other dtypes, and on
  BOOL it is computed with autocompute or refused with the ".new() first" hint without, as
  for Matrix.
- An infix expression forwards TensorExpression's methods (to_dense, ...), so with
  autocompute they compute on BOOL and raise for other dtypes, as Matrix's do.
- shape no longer needs the expression (the broadcast of the operands' shapes), so
  (X | Y).shape works for any dtype, as Matrix's nrows does.
- With an operator (binary.plus(X | Y), monoid.max(X & Y)) every dtype works as before.

The user guide's infix paragraph and the infix class's docstring say so, and decisions.md
records it as 27 (who, why, revisit if users ask for a default operator). The Tensor's own
code, oracles, notebooks, docs examples, benchmarks/tensor_ops.py and the benchmark kit use
| and & only on BOOL or with an operator; one test assigned (B | B) of FP64 to a region and
now checks the refusal and assigns binary.plus(B | B).

A new test checks, for FP64 and INT64, | and & refused with Matrix's message through
.new(), C <<, a region, ewise_add, +, einsum and to_dense, with autocompute off and on,
the same refusal from Matrix, shape still right, binary.plus / binary.times results against
numpy, and BOOL unchanged against Matrix.
…nse table

A broadcast ewise_mult that puts the small side on the nonempty rows of a hypersparse row
group (_bcast_mult's "diagonal rows, nonempty rows" option, _expand_on_rows) looked the
small side's values up in two dense numpy tables as long as its Vector: np.zeros(v._size)
for the values and for a presence mask. That is sized by the small side's axes, not by its
entries: for an axis of 2^40 it asks for 9 TiB. CI's Linux and Windows jobs failed so in
test_broadcast_mult_past_2_60_no_flip (Y.ewise_mult(X, pow), Y of shape (1, 2^40, 1, 1));
macOS hands out zeroed memory lazily, so it passed there.

The values are now found by a binary search of the small side's sorted indices
(tl.lookup_sorted, new): O((entries + nonempty rows) log entries), with the same results,
dtype included, and positions without a value left out as before.

The new test (test_broadcast_on_nonempty_rows_sized_by_entries) bounds the peak with
tracemalloc, which counts numpy's lazily zeroed allocations too: before the fix the peak
was 9895604665746 bytes on macOS (the test fails), after it is under 2^24.
A sweep of the Tensor modules for the hazard 5fb13f9 fixed (dense numpy allocations sized
by an axis, a group, a Vector or a Matrix dimension rather than by entries), with numpy's
constructors made to refuse large sizes while operations ran on Tensors of 2^40 axes and
past 2^60 positions, found four more that a huge but sparse operand reaches:

- A broadcast ewise_mult with the small side's axes inside a wider group expanded it
  through a selector over every position of the group (tl.selector: an arange of them).
  Past 2^60 positions in total that may be the one option left, over 2^60 columns: numpy
  refuses outright ("array is too big"). A group of more positions than max(4 * nnz, 4096)
  (_perm_matrix's rule) now takes the small side onto its nonempty rows or columns
  (_expand_on_ids; tl.nonempty_col_ids, new), as a hypersparse row group already did. The
  costs are as before; the census labels the option "diagonal rows/cols, nonempty
  rows/cols".
- A TTV along an axis inside a hypersparse row group filled a dense table and a presence
  mask over the whole axis (9 TiB for 2^40). rows_selector now takes w's sorted indices and
  values and finds them by tl.lookup_sorted.
- The Numba MTTKRP kernel's owner and column-block modes fill a dense (mode x R) output
  whatever the data (only per-thread buffers were bounded): 24 TiB for a mode of 2^40 at
  R = 3. It now declines past the per-thread bound, max(R * nnz, 2^22), and the GraphBLAS
  products, sized by the entries, take it.
- take_along_axis listed every point of the broadcast shape for the gather, so indices
  that broadcast over a long axis of the Tensor listed all of it (64 TiB per coordinate
  array for 2^40). Past max(4 * (nnz + indices), 4096) points it now joins the entries
  with the indices (_along_by_entries), sized by them and the result. It also computes the
  broadcast shape itself: numpy's broadcast_shapes refuses one past 2^63 positions, which
  read as indices that do not broadcast.

The other allocations the sweep found are sized by entries, by the operands' own storage
(a sparse Matrix's row pointers, full factors), by the user's index arrays, or are
selectors priced against an option sized by the entries that always exists; the report
lists each.

Each new test bounds the peak with tracemalloc (it counts numpy's lazily zeroed
allocations) and fails before this commit on macOS: the broadcast past 2^60 by numpy's
refusal; the TTV at 9895604658905 bytes; MTTKRP at 8858887824 for a mode of 2^28 (the test
runs 2^40 after it); take_along_axis at 411238565 for an axis of 2^20 (likewise). The
broadcast and take_along_axis tests also check the new ways against numpy at moderate
sizes, forced or past their threshold.
…ut a search where cheap

Round 5's sweep moved broadcasts over a group wider than max(4 * nnz, 4096)
positions onto its nonempty rows or columns, and looked the small side's
values up by a binary search instead of a dense table. At ordinary sizes
both cost time:

- The wide-group limit is now max(4 * nnz, 4096, the selector cache's
  bound of 2^25). Below it the selector is shape-only and cached, so a
  repeated broadcast pays one product over the group, where the nonempty
  columns are found again on every call. Past it (and past 2^60 positions,
  the case round 5 fixed) the nonempty rows or columns run as before.
- nonempty_col_ids sorts the column indices and masks the changes instead
  of np.unique: 22 against 146-153 ms for 1e6 indices.
- The nonempty-columns option is priced per entry (C_NONEMPTY_COLS, fitted
  forced on hyper and Zipf), not by the selector's per-position price; no
  kit decision changes (320 rows, census).
- lookup_sorted indexes a full vector directly and uses a dense table while
  the vector is within max(4 * (entries + positions), 4096); the binary
  search, whose memory follows the data, runs only past that.

Against 7ebce90, interleaved, 3 reps (load 3.3-5.9; ms, 8 / 1 threads):
NELL-2 mul_vec0 39.3 -> 36.0 / 174 -> 171, mul_vec1 39.3 -> 36.2 /
177 -> 174, ttv0 86.6 -> 83.6, ttv1 74.3 -> 70.9 (1 thread within 1%).
Forced broadcast over a column group of G positions, 1e6 entries (2 reps):
G = 8-32 x nnz, 156-170 -> 6.3-20.4 (8t), 259-315 -> 18-62 (1t), the cached
selector again; G = 128 x nnz, past the cache, 161-200 -> 56 and 269-323
-> 163, by the faster distinct columns. Results are the same.

Tests pin the choices: the forced wide-group cases keep the selector under
the cache's bound and take the nonempty ids with it lowered; lookup_sorted's
three ways are checked against each other and its search's memory with
tracemalloc. The round 5 tracemalloc tests still pass.
to_pydata_sparse's docstring and the user guide now say that a "gcxs"
export holds a row pointer with one slot per position of its compressed
axes, so its memory follows the product of their lengths whatever the
entries (two compressed axes of 2^20 make 2^40 slots), and to choose
compressed_axes to keep it small; "coo" (and "dok") follow the entries.
Documented rather than refused: any threshold would be arbitrary, and the
user asked for GCXS.
reshape_owned copied instead of reshaping in place past 2^31 rows or
columns because SuiteSparse 10.5.1's in-place reshape kept 32-bit
indices that cannot hold the new ones (Issue 462). 10.5.2 fixed that,
so on 10.5.2 and later it now always reshapes in place, which skips a
copy of the values; older versions still copy. Results do not change.

Checked against the library: in-place and copying reshapes of 8 shapes
crossing 2^31 or 2^32 (hypersparse and sparse, rowwise and columnwise
order) all agree with numpy on 10.5.2. On 10.5.1, 14 of the 32
in-place reshapes are wrong, and 10 columnwise copies are wrong or raise
InvalidValue, all with a dimension past 2^32.

test_reshape_owned_past_2_32 runs reshape_owned both ways where the
library allows it (copy only before 10.5.2): (4, 2^31) to (1, 2^33),
whose last entry 10.5.1 put at column 2^32 - 1.
… in its test

From the independent review of the SuiteSparse 10.5.2 reshape commit: the module comment said
2^32 where _RESHAPE_IN_PLACE_PAST_2_31 and reshape_owned use 2^31, and the test resized a Matrix
to the shape from_coo had already given it.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant