Skip to content

27. Tensor Contractions as Matrix Products

Mountain goats in space helmets on Mercury

What you will understand: what a tensor contraction is, why every contraction is "a matrix multiply in disguise", and how to take the disguise off: rearrange the axes of two tensors of any rank so that the entire contraction becomes a single matrix product, which is the one operation Mountain Goat is built around. The matrix product is Mountain Goat's compiled @; the axis bookkeeping around it is C++. The result is compared, bit for bit, with the definition of a contraction written out directly, on the worked examples of the Tensor Contractions (CPU) appendix of the companion book CUDA From First Principles, plus seven more cases.

What you need to know first: Chapter 21 (@, the matrix product), Chapter 20 (calling compiled Mountain Goat from C++), Chapter 24 (the loop-order option) and, for the final comparison, Chapter 14 (the run-time inner-dimension check). The appendix is not required; its ideas are explained here, and the page links to it.

Compile and run

cd docs/part24/code && ./build.sh                  # once: the newest compiler (Chapter 24's mg-opt)
cd ../../part27/code
cpp/run.sh                                         # compile contract.mg, build the C++ demo against it, run it (-> cpp/run_out.txt)
./contraction_mutation.sh > contraction_mutation_out.txt   # breaks the code on purpose, nine ways
cd ../../part15/code && ./run_lit.sh               # the whole test suite
Every listing and output on this page comes from these commands.

Why this page shows FAIL lines, and why that is good

The mutation output near the end comes from a negative control: the contraction code deliberately broken and run against the tests. There a "caught" line (a test failed) is the expected, wanted result. The baseline line must say the tests pass, and it does. The trap skip output in the middle of the page is also meant to abort: that is the check working.

Where this comes from

CUDA From First Principles (a companion book by the same author) has an appendix, Tensor Contractions, From First Principles (CPU) (source: docs/appendix/tensor-contractions-cpu.md in the MaximumBeings/cuda-from-first-principles repository), which defines a contraction, implements a generic N-dimensional contract() in C++ with nested index loops, verifies it against numpy.tensordot, shows what goes wrong without an axis check, and measures loop order. Its opening line is the idea of this chapter: "Every tensor contraction you will ever write is either a matrix multiply in disguise, or a small collection of matrix multiplies stapled together."

Mountain Goat is a language about two-dimensional matrices, so this chapter takes that sentence literally. It reimplements the appendix's definition (as an independent checker, written fresh for this chapter) and then does the contraction the other way, through a matrix product, which is how libraries do it. The appendix's two worked examples and its axis-mismatch case are reproduced exactly. (The environment this book is written in could not open the appendix's web page, so everything above comes from reading its source in the repository, at commit 6b3f95f.)

Primer: a contraction is a matrix product in disguise

A tensor is a block of numbers with some number of axes, each with a size: a vector has one axis, a matrix two, and a [2, 3, 4] tensor three. A contraction takes two tensors, pairs up some axes of the first with some axes of the second (the paired axes must have the same size), and for every combination of the other ("free") axes sums the products of matching elements over every value of the paired axes. The paired axes disappear from the result; the free axes remain, first A's and then B's.

The matrix product is the simplest case: C[i][j] = Σ_k A[i][k] · B[k][j] pairs A's axis 1 with B's axis 0. The disguise comes off with one observation. Group A's axes into two bundles, the free ones and the contracted ones, and think of each bundle as a single combined axis (a row number counting all combinations of the free axes, a column number counting all combinations of the contracted ones). Then A is just a matrix [M, K], where M is the product of A's free sizes and K is the product of the contracted sizes. Do the same for B with the contracted bundle first: a matrix [K, N]. The contraction is then exactly C[m][n] = Σ_k A'[m][k] · B'[k][n]: a matrix product. Reshape C back to A's free sizes followed by B's free sizes and you have the answer.

The recipe is called transpose-reshape-GEMM-reshape (GEMM is the traditional name for a general matrix multiply). Each word is a step:

  1. Transpose (more exactly, permute): reorder A's axes so the free ones come first and the contracted ones last; reorder B's so the contracted ones come first and the free ones last. The contracted axes must be listed in the same pairing order for both. This step copies data, unless the axes are already in that order.
  2. Reshape: a tensor stored row-major (last axis fastest) is already a flat buffer; viewing it as a matrix with M rows and K columns needs no copying at all, only a different way of reading the same numbers.
  3. GEMM: one matrix product [M, K] · [K, N]. This is where all the arithmetic is, and in this chapter it is one compiled Mountain Goat function.
  4. Reshape the [M, N] result back to the output shape (again, no copying).

A worked example. Contract A of shape [2, 3, 4] with B of shape [3, 4, 5], pairing A's axes {1, 2} with B's {0, 1}. A's free axis is {0} (size 2) and its contracted axes are sizes 3 and 4, so M = 2 and K = 3·4 = 12; B's contracted axes are already first (sizes 3 and 4) and its free axis has size 5, so N = 5. Nothing needs permuting: A is read as a 2×12 matrix and B as a 12×5 matrix, one product gives 2×5, which is the output shape. The multiply-add count is M·N·K = 2·5·12 = 120, which is the appendix's formula (output elements times the product of contracted sizes: 10 · 12).

Degenerate cases fall out for free. With no contracted axes the empty product gives K = 1: a [M, 1] times a [1, N] is the outer product. With every axis contracted, M = N = 1: a [1, K] times a [K, 1] is the dot product, a 1×1 matrix, which is the scalar (a tensor of rank 0, the empty shape []).

The Mountain Goat part

All of the arithmetic is this:

# The entire arithmetic of a tensor contraction, as one Mountain Goat function.
# Any contraction can be rearranged into a single matrix product (see tensor.h): the sum over the contracted axes becomes the sum over
# the inner dimension of the product. The row and column counts are dynamic (?), so one compiled function serves every contraction.
def gemm(a: tensor[?x?], b: tensor[?x?]) = a @ b

One function, gemm: a matrix product with dynamic row and column counts, so a single compiled function serves every contraction whatever its shape. (Mountain Goat has no tensors of rank 3 or more; that is why the reshaping lives in C++.) It is built with -O2 and the ikj loop order from Chapter 24.

The C++ part: the definition, and the matrix-product version

// Tensor contractions from first principles, with Mountain Goat doing the arithmetic.
//
//   contract_reference : the DEFINITION, written directly: for every output index, add up the products over every contracted index.
//                        No matrix product anywhere. It exists to check the other one.
//   contract_ttgt      : "transpose, reshape, GEMM, reshape": permute each tensor so its free axes come first (A) or last (B) and its
//                        contracted axes come last (A) or first (B), flatten each to a matrix, multiply the two matrices with the
//                        compiled Mountain Goat gemm (an @), and reshape the product to the output shape.
//
// Mountain Goat only has two-dimensional matrices, so the bookkeeping over N axes (permuting, flattening) is C++ here, and the whole sum
// of products is the language's matrix product.
#pragma once
#include <algorithm>
#include <cstdint>
#include <functional>
#include <stdexcept>
#include <string>
#include <vector>
#include "contract.h"   // generated by `mgc lib contract.mg`: mg::Matrix and mg::gemm

struct Tensor {
  std::vector<int> shape;       // row-major: the last axis varies fastest. An empty shape is a scalar (one element).
  std::vector<double> data;
  explicit Tensor(std::vector<int> s) : shape(std::move(s)) {
    long long total = 1; for (int d : shape) total *= d;
    data.assign((size_t)total, 0.0);
  }
};

inline long long product(const std::vector<int> &dims) { long long p = 1; for (int d : dims) p *= d; return p; }   // empty product is 1

inline std::vector<long long> row_major_strides(const std::vector<int> &shape) {
  std::vector<long long> strides(shape.size()); long long acc = 1;
  for (int i = (int)shape.size() - 1; i >= 0; i--) { strides[i] = acc; acc *= shape[i]; }
  return strides;
}

// Visit every multi-index of `shape` in row-major order (once, with an empty index, if the shape is empty).
inline void for_each_index(const std::vector<int> &shape, const std::function<void(const std::vector<int> &)> &visit) {
  std::vector<int> idx(shape.size(), 0);
  while (true) {
    visit(idx);
    int axis = (int)shape.size() - 1;
    while (axis >= 0 && ++idx[axis] == shape[axis]) { idx[axis] = 0; axis--; }
    if (axis < 0) return;
  }
}

// Check the axis lists. Throws std::invalid_argument; the first problem found is reported.
inline void validate(const Tensor &A, const std::vector<int> &axes_a, const Tensor &B, const std::vector<int> &axes_b) {
  if (axes_a.size() != axes_b.size()) throw std::invalid_argument("contract: axes_a and axes_b must have the same length");
  auto check_axes = [](const char *name, const Tensor &T, const std::vector<int> &axes) {
    std::vector<int> seen;
    for (int a : axes) {
      if (a < 0 || a >= (int)T.shape.size()) throw std::invalid_argument(std::string("contract: axis ") + std::to_string(a) + " is out of range for " + name);
      if (std::find(seen.begin(), seen.end(), a) != seen.end()) throw std::invalid_argument(std::string("contract: axis ") + std::to_string(a) + " is listed twice for " + name);
      seen.push_back(a);
    }
  };
  check_axes("A", A, axes_a); check_axes("B", B, axes_b);
  for (size_t k = 0; k < axes_a.size(); k++) {
    int da = A.shape[axes_a[k]], db = B.shape[axes_b[k]];
    if (da != db)
      throw std::invalid_argument("contract: mismatched dimension on contracted axis pair " + std::to_string(k) + " (A.shape[" + std::to_string(axes_a[k]) +
                                  "]=" + std::to_string(da) + " vs B.shape[" + std::to_string(axes_b[k]) + "]=" + std::to_string(db) + ")");
  }
}

inline std::vector<int> free_axes(int rank, const std::vector<int> &contracted) {
  std::vector<int> f; for (int a = 0; a < rank; a++) if (std::find(contracted.begin(), contracted.end(), a) == contracted.end()) f.push_back(a); return f;
}

// ---- 1. The definition, directly ----
inline Tensor contract_reference(const Tensor &A, const std::vector<int> &axes_a, const Tensor &B, const std::vector<int> &axes_b) {
  validate(A, axes_a, B, axes_b);
  std::vector<int> fa = free_axes((int)A.shape.size(), axes_a), fb = free_axes((int)B.shape.size(), axes_b);
  std::vector<int> out_shape, contracted_shape;
  for (int a : fa) out_shape.push_back(A.shape[a]);
  for (int b : fb) out_shape.push_back(B.shape[b]);
  for (int a : axes_a) contracted_shape.push_back(A.shape[a]);
  Tensor C(out_shape);
  auto sA = row_major_strides(A.shape), sB = row_major_strides(B.shape), sC = row_major_strides(out_shape);
  for_each_index(out_shape, [&](const std::vector<int> &out_idx) {
    long long offA = 0, offB = 0, offC = 0;
    for (size_t i = 0; i < fa.size(); i++) offA += out_idx[i] * sA[fa[i]];
    for (size_t j = 0; j < fb.size(); j++) offB += out_idx[fa.size() + j] * sB[fb[j]];
    for (size_t i = 0; i < out_idx.size(); i++) offC += out_idx[i] * sC[i];
    double acc = 0.0;
    for_each_index(contracted_shape, [&](const std::vector<int> &c) {
      long long a = offA, b = offB;
      for (size_t k = 0; k < c.size(); k++) { a += c[k] * sA[axes_a[k]]; b += c[k] * sB[axes_b[k]]; }
      acc += A.data[a] * B.data[b];
    });
    C.data[offC] = acc;
  });
  return C;
}

// ---- 2. Transpose, reshape, GEMM, reshape ----
// Copy T into a new tensor whose axis i is T's axis perm[i].
inline Tensor permute(const Tensor &T, const std::vector<int> &perm) {
  std::vector<int> shape; for (int p : perm) shape.push_back(T.shape[p]);
  Tensor R(shape);
  auto sT = row_major_strides(T.shape); long long pos = 0;
  for_each_index(shape, [&](const std::vector<int> &idx) {
    long long off = 0; for (size_t i = 0; i < idx.size(); i++) off += idx[i] * sT[perm[i]];
    R.data[pos++] = T.data[off];
  });
  return R;
}

// With validate_axes = false the axis check is skipped (to show what Mountain Goat's own run-time check does about a mismatch).
inline Tensor contract_ttgt(const Tensor &A, const std::vector<int> &axes_a, const Tensor &B, const std::vector<int> &axes_b, bool validate_axes = true) {
  if (validate_axes) validate(A, axes_a, B, axes_b);
  std::vector<int> fa = free_axes((int)A.shape.size(), axes_a), fb = free_axes((int)B.shape.size(), axes_b);
  // A: free axes first, contracted axes last.   B: contracted axes first (in the same pairing order), free axes last.
  std::vector<int> perm_a = fa; perm_a.insert(perm_a.end(), axes_a.begin(), axes_a.end());
  std::vector<int> perm_b = axes_b; perm_b.insert(perm_b.end(), fb.begin(), fb.end());
  Tensor Ap = permute(A, perm_a), Bp = permute(B, perm_b);
  std::vector<int> free_a_dims, free_b_dims, contracted_a, contracted_b;
  for (int a : fa) free_a_dims.push_back(A.shape[a]);
  for (int b : fb) free_b_dims.push_back(B.shape[b]);
  for (int a : axes_a) contracted_a.push_back(A.shape[a]);
  for (int b : axes_b) contracted_b.push_back(B.shape[b]);
  // Reshape = reinterpret the (now contiguous) buffer as a matrix: no data moves.
  mg::Matrix Am((int64_t)product(free_a_dims), (int64_t)product(contracted_a), Ap.data);
  mg::Matrix Bm((int64_t)product(contracted_b), (int64_t)product(free_b_dims), Bp.data);
  mg::Matrix Cm = mg::gemm(Am, Bm);                      // <-- the whole sum of products, in compiled Mountain Goat
  std::vector<int> out_shape = free_a_dims; out_shape.insert(out_shape.end(), free_b_dims.begin(), free_b_dims.end());
  Tensor C(out_shape);
  C.data = Cm.data;                                       // reshape back: again just reinterpretation
  return C;
}

Two functions with the same signature. contract_reference is the definition, written directly: for each output index, add up A·B over every combination of the contracted indices, with no matrix product anywhere. It exists only to check the other. contract_ttgt is the recipe: it works out the free and contracted axes, builds the two permutations (perm_a is the free axes followed by the contracted ones, perm_b the contracted ones followed by the free ones), calls permute (the only step that moves data), wraps each result as an mg::Matrix whose sizes are the products of the bundled axes, calls mg::gemm, and reads the product back as a tensor of the output shape. validate checks the axis lists before anything else: equal length, every axis in range and not repeated, and every contracted pair the same size.

The demo program runs the cases and prints, for each, whether the two versions agree bit for bit:

// Tensor contractions, computed two ways and compared: the definition (contract_reference) and a Mountain Goat matrix product (contract_ttgt).
// The first two cases are the worked examples of the "Tensor Contractions (CPU)" appendix of "CUDA From First Principles".
#include <cmath>
#include <cstdint>
#include <cstdio>
#include <cstring>
#include "tensor.h"

static bool same_bits(const Tensor &x, const Tensor &y) {
  return x.shape == y.shape && x.data.size() == y.data.size() && std::memcmp(x.data.data(), y.data.data(), x.data.size() * sizeof(double)) == 0;
}
static std::string shape_str(const std::vector<int> &s) { std::string r = "["; for (size_t i = 0; i < s.size(); i++) r += (i ? "," : "") + std::to_string(s[i]); return r + "]"; }
static void print_matrix(const Tensor &C) {
  for (int i = 0; i < C.shape[0]; i++) { std::printf("   "); for (int j = 0; j < C.shape[1]; j++) std::printf(" %8.2f", C.data[i * C.shape[1] + j]); std::printf("\n"); }
}

// Deterministic, deliberately NOT nice numbers (so a different order of additions would change the last bits).
static Tensor filled(std::vector<int> shape, int seed) {
  Tensor T(std::move(shape));
  for (size_t i = 0; i < T.data.size(); i++) T.data[i] = 0.1 * (double)((i * 7 + seed) % 13) - 0.45 + 1.0 / (double)(i + seed + 3);
  return T;
}

static int failures = 0;
static void check(const char *name, const Tensor &A, const std::vector<int> &aa, const Tensor &B, const std::vector<int> &ab) {
  Tensor ref = contract_reference(A, aa, B, ab), got = contract_ttgt(A, aa, B, ab);
  bool ok = same_bits(ref, got);
  if (!ok) failures++;
  std::printf("  %-42s A%-9s B%-9s -> %-9s matches the definition bit for bit: %s\n", name, shape_str(A.shape).c_str(), shape_str(B.shape).c_str(), shape_str(got.shape).c_str(), ok ? "yes" : "NO");
}

int main(int argc, char **argv) {
  if (argc > 1 && std::strcmp(argv[1], "trap") == 0) {
    // The appendix's trap: contract a [3,2] with a [4,5] over axes of different sizes (2 and 4).
    Tensor A = filled({3, 2}, 1), B = filled({4, 5}, 2);
    std::printf("[3,2] x [4,5] over axis 1 of A (size 2) and axis 0 of B (size 4), with the axis check %s:\n", argc > 2 ? "SKIPPED" : "on");
    if (argc > 2) { std::fflush(stdout); contract_ttgt(A, {1}, B, {0}, /*validate_axes=*/false); std::printf("not reached\n"); return 0; }
    try { contract_ttgt(A, {1}, B, {0}); } catch (const std::invalid_argument &e) { std::printf("  caught std::invalid_argument: \"%s\"\n", e.what()); }
    return 0;
  }

  std::printf("1. Matrix multiply, as contract(A, axis 1, B, axis 0)   (appendix H.3)\n");
  {
    Tensor A({3, 2}), B({2, 4});
    double a[] = {1, 2, 3, 4, 5, 6}, b[] = {1, 0, 2, 1, 0, 1, 1, 2};
    std::copy(a, a + 6, A.data.begin()); std::copy(b, b + 8, B.data.begin());
    Tensor C = contract_ttgt(A, {1}, B, {0});
    std::printf("   C =\n"); print_matrix(C);
    double want[] = {1, 2, 4, 5, 3, 4, 10, 11, 5, 6, 16, 17};
    bool ok = C.shape == std::vector<int>{3, 4} && std::equal(want, want + 12, C.data.begin());
    std::printf("   equals the appendix's numpy-checked values: %s\n", ok ? "yes" : "NO"); if (!ok) failures++;
  }

  std::printf("\n2. A double contraction: [2,3,4] x [3,4,5] over A's axes {1,2} and B's axes {0,1}   (appendix H.4)\n");
  {
    Tensor A({2, 3, 4}), B({3, 4, 5});
    for (size_t i = 0; i < A.data.size(); i++) A.data[i] = (double)(i % 7) - 3.0;
    for (size_t i = 0; i < B.data.size(); i++) B.data[i] = (double)(i % 5) - 2.0;
    Tensor C = contract_ttgt(A, {1, 2}, B, {0, 1});
    {
      // The reshape: A's free axis {0} against its contracted axes {1,2}, and B's contracted axes {0,1} against its free axis {2}.
      long long m = A.shape[0], k = (long long)A.shape[1] * A.shape[2], n = B.shape[2];
      std::printf("   as matrices: %lldx%lld times %lldx%lld  ->  %lld multiply-adds (the appendix's formula T x product of contracted sizes: %lld x %lld = %lld)\n",
                  m, k, k, n, m * n * k, m * n, (long long)B.shape[0] * B.shape[1], m * n * k);
    }
    std::printf("   output shape %s\n   C =\n", shape_str(C.shape).c_str()); print_matrix(C);
    double want[] = {10, 5, 0, -5, -10, 2, 1, 0, -1, -2};
    bool ok = C.shape == std::vector<int>{2, 5} && std::equal(want, want + 10, C.data.begin());
    std::printf("   equals the appendix's numpy.tensordot values: %s\n", ok ? "yes" : "NO"); if (!ok) failures++;
  }

  std::printf("\n3. More contractions, each compared with the definition (bit for bit, on numbers that are not 'nice')\n");
  check("matrix product", filled({3, 2}, 1), {1}, filled({2, 4}, 2), {0});
  check("double contraction", filled({2, 3, 4}, 1), {1, 2}, filled({3, 4, 5}, 2), {0, 1});
  check("contracted axis in the FRONT of A", filled({2, 3, 4}, 3), {0}, filled({2, 5}, 4), {0});
  check("contracted axes listed in reverse order", filled({2, 3, 4}, 5), {2, 1}, filled({4, 3, 5}, 6), {0, 1});
  check("contracted axes in the middle of both", filled({5, 2, 3}, 7), {1, 2}, filled({4, 3, 2, 6}, 8), {2, 1});
  check("vector times matrix", filled({4}, 9), {0}, filled({4, 6}, 10), {0});
  check("dot product (every axis contracted)", filled({2, 3, 4}, 11), {0, 1, 2}, filled({2, 3, 4}, 12), {0, 1, 2});
  check("outer product (no contracted axis)", filled({3}, 13), {}, filled({4}, 14), {});
  check("rank 3 x rank 3, one shared axis", filled({3, 4, 5}, 15), {2}, filled({5, 2, 3}, 16), {0});

  std::printf("\n4. The axis check (appendix H.5): [3,2] x [4,5] over axes of size 2 and 4\n");
  {
    Tensor A = filled({3, 2}, 1), B = filled({4, 5}, 2);
    try { contract_ttgt(A, {1}, B, {0}); std::printf("   NOT rejected: NO\n"); failures++; }
    catch (const std::invalid_argument &e) { std::printf("   rejected, as it should be: \"%s\"\n", e.what()); }
    for (auto bad : std::vector<std::pair<std::vector<int>, std::vector<int>>>{{{1}, {0, 1}}, {{5}, {0}}, {{1, 1}, {0, 1}}}) {
      try { contract_ttgt(A, bad.first, B, bad.second); std::printf("   bad axis list NOT rejected: NO\n"); failures++; }
      catch (const std::invalid_argument &e) { std::printf("   also rejected: \"%s\"\n", e.what()); }
    }
  }

  std::printf("\n%s\n", failures == 0 ? "All checks passed." : "SOME CHECKS FAILED.");
  return failures == 0 ? 0 : 1;
}

The run

#!/bin/sh
# Compile contract.mg (optimized, ikj loop order), build the C++ demo against it and run it; then run the two "trap" modes. Output: run_out.txt
HERE=$(cd "$(dirname "$0")" && pwd); W=$HERE/../work/contract; mkdir -p $W
$HERE/../../../part24/code/mgc lib $HERE/contract.mg -o $W -O2 --matmul-order ikj > /dev/null
cp $W/contract.h $HERE/contract_generated.h
clang++-18 -std=c++17 -O2 -Wall -Wextra -I$W -I$HERE $HERE/contract_demo.cpp $W/contract.o -o $W/contract_demo
echo "--- ./contract_demo"; $W/contract_demo; echo "(exit status $?)"
echo; echo "--- ./contract_demo trap          (the axis check on)"; $W/contract_demo trap
echo; echo "--- ./contract_demo trap skip     (the axis check skipped: what does Mountain Goat's own run-time check say?)"
$W/contract_demo trap skip; echo "(exit status $?)"

--- ./contract_demo
1. Matrix multiply, as contract(A, axis 1, B, axis 0)   (appendix H.3)
   C =
        1.00     2.00     4.00     5.00
        3.00     4.00    10.00    11.00
        5.00     6.00    16.00    17.00
   equals the appendix's numpy-checked values: yes

2. A double contraction: [2,3,4] x [3,4,5] over A's axes {1,2} and B's axes {0,1}   (appendix H.4)
   as matrices: 2x12 times 12x5  ->  120 multiply-adds (the appendix's formula T x product of contracted sizes: 10 x 12 = 120)
   output shape [2,5]
   C =
       10.00     5.00     0.00    -5.00   -10.00
        2.00     1.00     0.00    -1.00    -2.00
   equals the appendix's numpy.tensordot values: yes

3. More contractions, each compared with the definition (bit for bit, on numbers that are not 'nice')
  matrix product                             A[3,2]     B[2,4]     -> [3,4]     matches the definition bit for bit: yes
  double contraction                         A[2,3,4]   B[3,4,5]   -> [2,5]     matches the definition bit for bit: yes
  contracted axis in the FRONT of A          A[2,3,4]   B[2,5]     -> [3,4,5]   matches the definition bit for bit: yes
  contracted axes listed in reverse order    A[2,3,4]   B[4,3,5]   -> [2,5]     matches the definition bit for bit: yes
  contracted axes in the middle of both      A[5,2,3]   B[4,3,2,6] -> [5,4,6]   matches the definition bit for bit: yes
  vector times matrix                        A[4]       B[4,6]     -> [6]       matches the definition bit for bit: yes
  dot product (every axis contracted)        A[2,3,4]   B[2,3,4]   -> []        matches the definition bit for bit: yes
  outer product (no contracted axis)         A[3]       B[4]       -> [3,4]     matches the definition bit for bit: yes
  rank 3 x rank 3, one shared axis           A[3,4,5]   B[5,2,3]   -> [3,4,2,3] matches the definition bit for bit: yes

4. The axis check (appendix H.5): [3,2] x [4,5] over axes of size 2 and 4
   rejected, as it should be: "contract: mismatched dimension on contracted axis pair 0 (A.shape[1]=2 vs B.shape[0]=4)"
   also rejected: "contract: axes_a and axes_b must have the same length"
   also rejected: "contract: axis 5 is out of range for A"
   also rejected: "contract: axis 1 is listed twice for A"

All checks passed.
(exit status 0)

--- ./contract_demo trap          (the axis check on)
[3,2] x [4,5] over axis 1 of A (size 2) and axis 0 of B (size 4), with the axis check on:
  caught std::invalid_argument: "contract: mismatched dimension on contracted axis pair 0 (A.shape[1]=2 vs B.shape[0]=4)"

--- ./contract_demo trap skip     (the axis check skipped: what does Mountain Goat's own run-time check say?)
[3,2] x [4,5] over axis 1 of A (size 2) and axis 0 of B (size 4), with the axis check SKIPPED:
(exit status 134)

Reading it.

  • Cases 1 and 2 are the appendix's worked examples. The matrix product [3,2]·[2,4] gives the values the appendix checked against numpy's @. The double contraction [2,3,4] × [3,4,5] over two axis pairs gives [[10, 5, 0, −5, −10], [2, 1, 0, −1, −2]], the values the appendix checked against numpy.tensordot. Neither check here uses numpy: the expected values come from the appendix's page, and the independent check is the definition.
  • Case 3 is nine contractions chosen to exercise the bookkeeping: a contracted axis at the front of A (which forces a real permutation), contracted axes listed in reverse order (so the pairing order matters), contracted axes in the middle of both tensors, a vector times a matrix, a dot product (the [] shape: a scalar), an outer product (no contracted axes), and rank-3 times rank-3 giving a rank-4 result. For all nine the matrix-product version equals the definition bit for bit, on numbers that are deliberately not "nice" (they are not exactly representable, so a different order of additions would change the last bits). That is expected for a reason worth stating: the flattened contracted index runs through the contracted axes in the same row-major order the definition uses, so each output element adds its products in the same order.
  • Case 4 is the appendix's section H.5, with extra axis mistakes: a mismatched pair (A.shape[1]=2 against B.shape[0]=4), lists of different lengths, an out-of-range axis, and a repeated axis. Each is rejected with a specific message. (The mismatch message uses the appendix's wording, so the two outputs line up.)

The trap, revisited: what happens if the check is skipped?

The appendix demonstrates why the axis check matters. An unchecked contraction, given the same mismatched [3,2] and [4,5], "does not crash and does not fail loudly: it silently sizes the contracted-index loop from A's axis alone (size 2) and reads only the first 2 of B's 4 rows", producing a plausible, wrong 3×5 result (that experiment is the appendix's, not reproduced here). The two trap runs above show the Mountain Goat version of the same experiment. With the check on, the mismatch is an exception with a precise message. With the check skipped (trap skip), the two tensors are flattened to a 3×2 and a 4×5 matrix and handed to gemm, and Mountain Goat's own run-time shape check (Chapter 14) stops the program: mg.matmul: inner dimensions differ at runtime on stderr, exit status 134, and nothing printed after it. So here the failure is not silent even if the C++ validation is deleted: the compiled matrix product refuses to multiply matrices whose inner sizes disagree, which is precisely the protection the appendix had to add by hand.

Loop order, again

The appendix's last section measures how the loop order of a matrix multiply changes its running time without changing its answer, on a plain C++ triple loop (ijk against ikj). That is the same experiment as Chapters 23 and 24, which ran it on Mountain Goat's own compiled @ and found the same effect (ikj several times faster at larger sizes on this machine, with identical results), and added --matmul-order ikj so the faster order can be chosen. A contraction done through gemm inherits that: the contract.mg here is built with ikj, and the test builds it both ways and requires identical output. No new timing is claimed in this chapter; the numbers in Chapters 23 and 24 apply to square products on one shared machine, not to these small demonstration shapes.

Tests

Two new lit tests in test/contraction/ (the suite now has 105 tests):

Test What it checks
tensor-contraction the demo, built twice (optimized with ikj, and with the plain defaults), must print every expected value and yes, with identical complete output from both builds; checks the appendix's two worked results, the 2x12 times 12x5 reshape line, all nine cases, and the exact message of each axis check
mismatch-trap with the check on, the exception message; with the check skipped, an abort with Mountain Goat's message on stderr and nothing after the "SKIPPED" line

Are the tests good enough? Break the code nine ways.

contraction_mutation.sh damages tensor.h or contract.mg one way at a time and re-runs the two tests:

#!/bin/sh
# READ THIS FIRST: this script edits the contraction code (cpp/tensor.h and cpp/contract.mg) to be WRONG, one way at a time, and re-runs
#   the Chapter 27 tests (test/contraction). A "caught" line (a test failed) is the EXPECTED, wanted result: it shows the tests can notice
#   that mistake. "NOT CAUGHT" would be a test gap. Every file is restored afterwards. Output: contraction_mutation_out.txt
HERE=$(cd "$(dirname "$0")" && pwd); D=$HERE/../..; LIT=$D/part15/code/run_lit.sh; W=$HERE/work/mut; mkdir -p $W
H=$HERE/cpp/tensor.h; G=$HERE/cpp/contract.mg
cp $H $W/tensor.orig; cp $G $W/contract.orig
echo "NOTE: this script deliberately breaks the contraction code. 'caught' lines are EXPECTED: they show the tests can detect the mistake."
run() { out=$(MG_TEST_TMP=$W/out sh $LIT --filter ':: contraction/' 2>&1); fails=$(echo "$out" | grep '^FAIL' | sed 's/.*:: //; s/ (.*//' | tr '\n' ' ')
  if [ -n "$fails" ]; then echo "$1: caught by $fails"; else echo "$1: NOT CAUGHT (a test gap)"; fi; }
mutate() {  # label file sed-expr
  cp $W/tensor.orig $H; cp $W/contract.orig $G; sed -i "$3" "$2"
  if cmp -s $H $W/tensor.orig && cmp -s $G $W/contract.orig; then echo "$1: MUTATION DID NOT APPLY"; else run "$1"; fi
  cp $W/tensor.orig $H; cp $W/contract.orig $G
}
out=$(MG_TEST_TMP=$W/out sh $LIT --filter ':: contraction/' 2>&1); echo "baseline (nothing broken): $(echo "$out" | grep -q '^FAIL' && echo 'UNEXPECTED FAILURE' || echo 'tests pass (expected)')"
mutate "A: contracted axes first, free axes last (swapped)"        $H 's/std::vector<int> perm_a = fa; perm_a.insert(perm_a.end(), axes_a.begin(), axes_a.end());/std::vector<int> perm_a = axes_a; perm_a.insert(perm_a.end(), fa.begin(), fa.end());/'
mutate "B: free axes first, contracted axes last (swapped)"        $H 's/std::vector<int> perm_b = axes_b; perm_b.insert(perm_b.end(), fb.begin(), fb.end());/std::vector<int> perm_b = fb; perm_b.insert(perm_b.end(), axes_b.begin(), axes_b.end());/'
mutate "B's contracted axes taken in increasing order, not paired" $H 's/std::vector<int> perm_b = axes_b;/std::vector<int> perm_b = axes_b; std::sort(perm_b.begin(), perm_b.end());/'
mutate "output shape lists the free axes of B first"                $H 's/std::vector<int> out_shape = free_a_dims; out_shape.insert(out_shape.end(), free_b_dims.begin(), free_b_dims.end());/std::vector<int> out_shape = free_b_dims; out_shape.insert(out_shape.end(), free_a_dims.begin(), free_a_dims.end());/'
mutate "the size check of a contracted pair is removed"           $H 's/if (da != db)$/if (false)/'
mutate "a repeated axis is no longer rejected"                    $H 's/if (std::find(seen.begin(), seen.end(), a) != seen.end()) throw/if (false) throw/'
mutate "an out-of-range axis is no longer rejected"               $H 's/if (a < 0 || a >= (int)T.shape.size()) throw/if (false) throw/'
mutate "the Mountain Goat gemm adds the product twice"            $G 's/= a @ b$/= a @ b + a @ b/'
mutate "the Mountain Goat gemm multiplies by the transpose"       $G 's/= a @ b$/= a @ transpose(b)/'
echo "restored: $(cmp -s $H $W/tensor.orig && cmp -s $G $W/contract.orig && echo 'all files are back to their originals' || echo 'RESTORE FAILED')"

NOTE: this script deliberately breaks the contraction code. 'caught' lines are EXPECTED: they show the tests can detect the mistake.
baseline (nothing broken): tests pass (expected)
A: contracted axes first, free axes last (swapped): caught by contraction/tensor-contraction.mlir 
B: free axes first, contracted axes last (swapped): caught by contraction/tensor-contraction.mlir 
B's contracted axes taken in increasing order, not paired: caught by contraction/tensor-contraction.mlir 
output shape lists the free axes of B first: caught by contraction/tensor-contraction.mlir 
the size check of a contracted pair is removed: caught by contraction/mismatch-trap.mlir contraction/tensor-contraction.mlir 
a repeated axis is no longer rejected: caught by contraction/tensor-contraction.mlir 
an out-of-range axis is no longer rejected: caught by contraction/mismatch-trap.mlir contraction/tensor-contraction.mlir 
the Mountain Goat gemm adds the product twice: caught by contraction/tensor-contraction.mlir 
the Mountain Goat gemm multiplies by the transpose: caught by contraction/tensor-contraction.mlir 
restored: all files are back to their originals

All nine are caught, including the four that are about the axis bookkeeping (swapping which axes come first, taking B's contracted axes in sorted order instead of paired order, listing B's free axes first in the output), the three about validation (the size check, repeated axes, out-of-range axes), and the two about the Mountain Goat function itself. Two details are worth noticing. The mutation "B's contracted axes taken in increasing order, not paired" is caught only because case 3 includes a contraction with axes listed in reverse order; on the simpler cases, where axes are already in increasing order, it changes nothing. And the validation mutations are caught by looking at the exact message: a repeated axis, for example, would otherwise be rejected anyway by the size check further down (with a different message), so a test that only asked "was it rejected?" would not notice that the repeated-axis check had been removed.

The full suite:

Testing Time: 7.72s

Total Discovered Tests: 105
  Passed: 105 (100.00%)

Limits and what is not established

  • Rank-2 arithmetic only. Mountain Goat has no tensors of rank 3 or more, so the permuting and reshaping are C++. Only the matrix product is the language's. A language-level reshape and permute are not built here; Chapter 28 builds them, and writes the same contractions with no C++.
  • The permutation copies. When the axes are not already in the needed order (case 3, "axis at the front of A"), permute copies the whole tensor before the product. That cost was not measured, and a production library would avoid or fuse it. When the axes already line up (the appendix's double contraction), no data moves.
  • No batch axes. A contraction pairs axes and sums them away; operations like batched matrix multiply (an axis present in both inputs and the output) are einsum, not a contraction in this sense, and are not supported.
  • Small shapes only were tested, all double, all tiny (the largest has a few hundred elements). Nothing here says how it behaves on large tensors, and nothing was timed.
  • The reference is mine, not the appendix's code. contract_reference was written for this chapter; it agrees with the appendix's printed values (cases 1 and 2) and with the matrix-product version on all nine cases, but the appendix's own program was not run here.
  • Bit-for-bit equality is shown for these cases on this build. It follows from the order of additions being the same, which is an argument plus nine tests, not a proof for every shape; a different compiler setting that reorders floating-point additions (for example, fast-math flags, which this build does not use) would break it.
  • The GPU side is not covered. The appendix has a companion for CUDA; Mountain Goat's GPU path is compile-only (Chapters 10 to 12) and nothing here runs on a GPU.
  • Platform: x86-64 Linux, LLVM 18.1.3.

Reproducing

cd docs/part24/code && ./build.sh
cd ../../part27/code
cpp/run.sh > cpp/run_out.txt 2>&1
./contraction_mutation.sh > contraction_mutation_out.txt
cd ../../part15/code && ./run_lit.sh          # 105 tests

Chapter summary

  • A contraction pairs axes of two tensors and sums products over them. Bundling each tensor's free axes into one axis and its contracted axes into another turns it into an ordinary matrix product: permute, reshape, multiply, reshape.
  • With Mountain Goat that gives a one-line function, gemm, for all the arithmetic; the C++ around it only rearranges axes.
  • It agrees bit for bit with a direct implementation of the definition on the appendix's two worked examples and nine further cases, including a front-of-A axis, reversed axis order, a dot product (rank 0), an outer product, and a rank-4 result.
  • Skipping the axis check does not silently corrupt the result: Mountain Goat's run-time inner-dimension check aborts with a message, where an unchecked hand-written loop (as the appendix shows) quietly drops half of one input.
  • The loop-order effect the appendix measures is the same one Chapters 23 and 24 measured and made selectable.
  • Nine injected bugs are all caught; two of them are caught only because the cases were chosen to exercise reversed axis order and exact error messages.

Self-check questions

Each answer is collapsed; try the question first.

  1. In one sentence each, what are free axes and contracted axes, and what does the output shape consist of?

    Answer

    Contracted axes are the paired axes that are summed over and disappear; free axes are all the others and survive. The output shape is A's free axes followed by B's free axes, each at its original size.

  2. A has shape [2,3,4] and B has shape [3,4,5]; the pairs are A's axes {1,2} with B's {0,1}. What matrices does the matrix-product version multiply, and how many multiply-adds does that take?

    Answer

    A is read as [2, 12] (the free axis, then the two contracted axes bundled: 3·4 = 12) and B as [12, 5] (the contracted axes bundled, then the free axis). The product is [2, 5] and takes 2·5·12 = 120 multiply-adds, equal to the appendix's formula: output elements (10) times the product of the contracted sizes (12).

  3. Why does the "reshape" step cost nothing in this case, and when does the "transpose" step have to copy data?

    Answer

    A row-major tensor is already a flat buffer, so viewing it as a matrix just changes how the same numbers are read. Here A's free axis is first and its contracted axes are last, and B's contracted axes are first and its free axis last, so no reordering is needed. If an axis is somewhere else (say, a contracted axis at the front of A), the axes must be reordered, which means copying the data into the new order.

  4. What are the matrix shapes for an outer product (no contracted axes) and for a dot product (all axes contracted), and why?

    Answer

    An outer product has K = 1, because the product of no sizes is 1: a [M, 1] times a [1, N] gives [M, N]. A dot product has no free axes, so M = N = 1 (again the empty product): a [1, K] times a [K, 1] gives [1, 1], the scalar, which the tensor code represents with the empty shape.

  5. Why does case 3 include a contraction with the contracted axes listed in reverse order, and what mutation does it catch?

    Answer

    The pairing order of the contracted axes matters: the axis listed first for A must meet the axis listed first for B. If B's contracted axes were instead taken in sorted order, simple cases where the axes are already increasing would still be right, and only a case with a reversed order would produce wrong numbers. The mutation "B's contracted axes taken in increasing order, not paired" is caught by exactly that case.

  6. Why are the results bit-for-bit identical to the definition rather than merely close?

    Answer

    For each output element both versions add the products in the same order: the flattened contracted index runs over the contracted axes in row-major order, the same order the direct definition's inner loop uses. Floating-point addition is not associative, but the same sequence of additions gives the same bits. The test uses numbers that are not exactly representable so that a different order would show.

  7. What happens if the axis validation is skipped and the contracted sizes differ, in the appendix's unchecked version and here?

    Answer

    In the appendix's unchecked version the loop is sized from one tensor's axis and quietly reads only part of the other tensor, giving a plausible but wrong result with no error. Here the two tensors are flattened to matrices with different inner sizes and handed to the compiled matrix product, whose own run-time check (Chapter 14) aborts with mg.matmul: inner dimensions differ at runtime. The failure is loud even without the C++ validation.

  8. Why does the mutation script check the exact error message rather than only that an exception was thrown?

    Answer

    Several mistakes are caught by more than one check. If the repeated-axis check were deleted, a later check (the size comparison) might still throw for the same input, with a different message; a test that only asked "was an exception thrown?" would pass, and the deleted check would go unnoticed. Checking the exact message distinguishes which check fired.

  9. Why can't this chapter claim anything about the speed of contractions?

    Answer

    Nothing was timed: the shapes are tiny and meant for checking correctness, and the permutation's copy cost was not measured. Chapters 23 and 24 measured square matrix products on one shared machine, and those numbers cannot be transferred to other shapes or to the permute-and-reshape overhead.

  10. What would be needed for Mountain Goat itself, not C++, to express a general contraction?

    Answer

    Tensors of rank greater than two, a reshape operation (a no-copy reinterpretation of a contiguous tensor as another shape) and a general permute (a transpose over any axes), each with a verifier, a lowering and front-end syntax and tests, as in Chapters 21 and 22. Alternatively a single contraction operation with axis-list attributes, lowered to the same permute, reshape and matrix product. The C++ code in this chapter is a specification of what those operations would have to do.