22. Broadcasting, Reductions and relu: Enough Operations for a Neural-Network Layer

What you will understand: how to express "add this row to every row", "total each column" and "clamp negatives to zero", the three ideas (broadcasting, reduction, an activation function) that turn matrix arithmetic into a neural-network layer. You will see how each becomes loops, why broadcasting is a rule about shapes and not new arithmetic, what a reduction does to a shape, and why a test suite can pass while a real bug hides in it.

What you need to know first: the language tour (syntax and shapes), Chapter 21 (the elementwise operations and @), and Chapters 13 and 14 for what ? means and what is checked at run time.

Compile and run

cd docs/part22/code
./build.sh                                  # builds ./build/mg-opt (Chapter 22's)
./mgc run examples/07_mlp_layer.mg          # compile one example to native code and run it
./mgc ptx examples/10_gpu_layer.mg          # compile one example to PTX (not run: no GPU here)
./demo.sh > demo_out.txt                    # all the examples; the output on this page is this file
./show.sh                                   # the lowered MLIR shown below  -> show/
cpp/run.sh                                  # build and run the C++ program
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

Parts of this page come from negative controls: the same tests run against an older build, or against deliberately broken code (a "mutation"). In those runs a FAIL is the expected, wanted result: it means a test noticed the problem, which is how we know the tests are worth anything. What would be wrong is the opposite, a deliberately broken build that passes everything. The baseline (the unmodified current build) must always show every test passing, and it does. Each script that does this says so in its own header and prints a reminder at the top of its output.

Primer: three ideas

Broadcasting. [[1, 2, 3], [4, 5, 6]] + [[10, 20, 30]] adds a 1x3 row to a 2x3 matrix. The shapes differ, so by Chapter 21's rule the program would be an error. But the intent is clear and extremely common (adding a bias to every row). Broadcasting is the rule that makes it legal: a dimension of size 1 is conceptually repeated to match the other operand, so the 1x3 row behaves as if it were two copies of itself stacked into a 2x3 matrix. Broadcasting is not a new arithmetic operation; it only changes shapes before the operation runs. The rule, per dimension: the sizes must be equal, or one of them must be 1 (in which case it grows to the other).

Reductions. A reduction combines many numbers into one: the sum of a row, the largest of a column. For a matrix there are two natural directions. This book names them by what you get: row_sum(a) gives one value per row (an M by 1 result, each row collapsed to a single number), and col_sum(a) gives one value per column (a 1 by N result). The same pair exists for max and for mean. The reduced dimension becomes size 1, which is exactly what lets the result be broadcast back against the original: a - col_mean(a) subtracts each column's average from that column.

Activation functions. A neural-network layer computes x @ w + bias and then applies a nonlinear function elementwise. The simplest and most common is relu (rectified linear unit): relu(x) = max(x, 0), which turns every negative number into zero and leaves the rest unchanged. Without it, stacking layers would collapse into one big matrix product.

Put together, relu(x @ w + bias) is one layer, and it uses everything this chapter adds. That is the target.

Design decisions

• Names say what you get. row_sum returns one value per row. A name like sum(a, axis=1) makes you remember which axis is which; these do not.
• Broadcasting is implicit and static. A size-1 dimension is broadcast automatically by + - * /, and the front end inserts an explicit mg.broadcast operation to do the repetition, so every later stage sees plain equal-shape arithmetic. It works only when the shapes involved are static (known numbers). A ? dimension is never broadcast: its size is unknown until the program runs, so it must be equal to the other operand's size at run time, exactly as in Chapter 14.
• Mean needs a count. row_mean is a sum divided by the number of columns. That number must be known when compiling, so row_mean on a ?-wide matrix is a compile error (the other reductions work on ?).
• relu propagates NaN. relu lowers to a floating-point maximum that returns NaN if the input is NaN, rather than hiding it as 0. That matches how most ML frameworks treat it, and it is what NVIDIA's max.NaN instruction does.

The operations (MgOps.td)

Three new definitions, in the same ODS language as Chapters 3 and 21:

def ReluOp : Mg_Op<"relu", [Pure]> {
  let summary = "elementwise max(x, 0), the rectified linear unit";
  let arguments = (ins F64Tensor:$input);
  let results = (outs F64Tensor:$result);
  let assemblyFormat = "$input attr-dict `:` type($input) `->` type($result)";
  let hasVerifier = 1;
}

def ReduceOp : Mg_Op<"reduce", [Pure]> {
  let summary = "reduce a matrix along one axis with sum or max";
  let description = [{
    `axis = 1` combines the elements of each row, giving an M x 1 result; `axis = 0` combines the
    elements of each column, giving a 1 x N result. `kind` is "sum" or "max".
  }];
  let arguments = (ins F64Tensor:$input, I64Attr:$axis, StrAttr:$kind);
  let results = (outs F64Tensor:$result);
  let assemblyFormat = "$input attr-dict `:` type($input) `->` type($result)";
  let hasVerifier = 1;
}

def BroadcastOp : Mg_Op<"broadcast", [Pure]> {
  let summary = "repeat a matrix along every dimension of size 1 to fill a larger static shape";
  let arguments = (ins F64Tensor:$input);
  let results = (outs F64Tensor:$result);
  let assemblyFormat = "$input attr-dict `:` type($input) `->` type($result)";
  let hasVerifier = 1;
}

mg.relu is a one-input elementwise op, like mg.neg. mg.reduce carries two attributes: axis (1 means "combine the elements of each row", giving an M by 1 result; 0 means "each column", giving 1 by N) and kind ("sum" or "max"). Mean is not a separate operation: the front end emits a sum followed by a divide by the (static) count. mg.broadcast has one operand and an explicit result type; its verifier says what shapes are allowed.

The verifiers (MgDialect.cpp)

mlir::LogicalResult ReluOp::verify() { return verifyElementwise(*this, {getInput().getType(), getResult().getType()}, "mg.relu"); }

mlir::LogicalResult ReduceOp::verify() {
  auto in = llvm::dyn_cast<mlir::RankedTensorType>(getInput().getType());
  auto out = llvm::dyn_cast<mlir::RankedTensorType>(getResult().getType());
  if (!in || !out || in.getRank() != 2 || out.getRank() != 2)
    return emitOpError("mg.reduce only supports rank-2 ranked tensors");
  if (getAxis() != 0 && getAxis() != 1)
    return emitOpError("mg.reduce: axis must be 0 or 1, got ") << getAxis();
  llvm::StringRef k = getKind();
  if (k != "sum" && k != "max")
    return emitOpError("mg.reduce: kind must be sum or max, got '") << k << "'";
  int64_t kept = 1 - getAxis();   // the dimension that survives
  if (!mg::dyn::compatibleDim(out.getDimSize(kept), in.getDimSize(kept)) || out.getDimSize(getAxis()) != 1)
    return emitOpError("mg.reduce: reducing axis ") << getAxis() << " of " << in << " must give a result with that axis of size 1, got " << out;
  return mlir::success();
}

mlir::LogicalResult BroadcastOp::verify() {
  auto in = llvm::dyn_cast<mlir::RankedTensorType>(getInput().getType());
  auto out = llvm::dyn_cast<mlir::RankedTensorType>(getResult().getType());
  if (!in || !out || in.getRank() != 2 || out.getRank() != 2)
    return emitOpError("mg.broadcast only supports rank-2 ranked tensors");
  if (!in.hasStaticShape() || !out.hasStaticShape())
    return emitOpError("mg.broadcast needs static shapes (dynamic broadcasting is not supported)");
  for (int d = 0; d < 2; ++d)
    if (in.getDimSize(d) != out.getDimSize(d) && in.getDimSize(d) != 1)
      return emitOpError("mg.broadcast: dimension ") << d << " of " << in << " can only grow from size 1, not to " << out;
  return mlir::success();
}

mg.relu reuses Chapter 21's shared elementwise check. mg.reduce checks the axis, the kind, and that the result collapsed exactly the reduced axis to size 1 and kept the other (accepting ?). mg.broadcast requires static shapes and allows an input dimension to differ from the result only if it is 1.

The lowerings (LowerToAffine.cpp)

relu is Chapter 21's negate with maximumf(x, 0) instead of negf:

struct ReluOpLowering : public OpConversionPattern<mg::ReluOp> {
  using OpConversionPattern::OpConversionPattern;
  LogicalResult matchAndRewrite(mg::ReluOp op, OpAdaptor adaptor,
                                ConversionPatternRewriter &rewriter) const override {
    auto loc = op.getLoc();
    auto tensorType = llvm::cast<RankedTensorType>(op.getResult().getType());
    auto memrefType = convertTensorToMemref(tensorType);
    Value in = adaptor.getInput();
    llvm::SmallVector<Value, 2> extents;
    for (int64_t d = 0; d < 2; ++d)
      extents.push_back(tensorType.isDynamicDim(d) ? mg::dyn::extentOf(rewriter, loc, in, d) : Value());
    auto alloc = mg::dyn::allocFor(rewriter, loc, memrefType, extents);
    Value zero = rewriter.create<arith::ConstantOp>(loc, rewriter.getF64FloatAttr(0.0));
    mg::dyn::nest2D(rewriter, loc, tensorType, extents, [&](OpBuilder &b, Location l, ValueRange ivs) {
      Value x = b.create<affine::AffineLoadOp>(l, in, ivs);
      b.create<affine::AffineStoreOp>(l, b.create<arith::MaximumFOp>(l, x, zero), alloc, ivs);
    });
    rewriter.replaceOp(op, alloc.getResult());
    return success();
  }
};

/// mg.reduce: initialise the result (0 for sum, -infinity for max), then one nest over the INPUT that
/// folds each element into result[i, 0] (axis 1) or result[0, j] (axis 0).

reduce is the new idea. It cannot be a single loop nest that stores each output once, because each output combines many inputs. Instead it runs two nests: the first writes the starting value into every output cell (0 for a sum, negative infinity for a max, because max(x, -infinity) is always x); the second walks the input and folds every element into the output cell it belongs to (result[i][0] when reducing a row, result[0][j] when reducing a column):

struct ReduceOpLowering : public OpConversionPattern<mg::ReduceOp> {
  using OpConversionPattern::OpConversionPattern;
  LogicalResult matchAndRewrite(mg::ReduceOp op, OpAdaptor adaptor,
                                ConversionPatternRewriter &rewriter) const override {
    auto loc = op.getLoc();
    auto inType = llvm::cast<RankedTensorType>(op.getInput().getType());
    auto outType = llvm::cast<RankedTensorType>(op.getResult().getType());
    auto memrefType = convertTensorToMemref(outType);
    Value in = adaptor.getInput();
    int64_t axis = op.getAxis();
    bool isMax = op.getKind() == "max";
    llvm::SmallVector<Value, 2> outExtents;
    for (int64_t d = 0; d < 2; ++d)
      outExtents.push_back(outType.isDynamicDim(d) ? mg::dyn::extentOf(rewriter, loc, in, d) : Value());
    auto alloc = mg::dyn::allocFor(rewriter, loc, memrefType, outExtents);
    double initVal = isMax ? -std::numeric_limits<double>::infinity() : 0.0;
    Value init = rewriter.create<arith::ConstantOp>(loc, rewriter.getF64FloatAttr(initVal));
    mg::dyn::nest2D(rewriter, loc, outType, outExtents, [&](OpBuilder &b, Location l, ValueRange ivs) {
      b.create<affine::AffineStoreOp>(l, init, alloc, ivs);
    });
    llvm::SmallVector<Value, 2> inExtents;
    for (int64_t d = 0; d < 2; ++d)
      inExtents.push_back(inType.isDynamicDim(d) ? mg::dyn::extentOf(rewriter, loc, in, d) : Value());
    mg::dyn::nest2D(rewriter, loc, inType, inExtents, [&](OpBuilder &b, Location l, ValueRange ivs) {
      Value zeroIdx = b.create<arith::ConstantIndexOp>(l, 0);
      Value i = ivs[0], j = ivs[1];
      llvm::SmallVector<Value, 2> dst = axis == 1 ? llvm::SmallVector<Value, 2>{i, zeroIdx} : llvm::SmallVector<Value, 2>{zeroIdx, j};
      Value acc = b.create<affine::AffineLoadOp>(l, alloc, dst);
      Value x = b.create<affine::AffineLoadOp>(l, in, ivs);
      Value r = isMax ? Value(b.create<arith::MaximumFOp>(l, acc, x)) : Value(b.create<arith::AddFOp>(l, acc, x));
      b.create<affine::AffineStoreOp>(l, r, alloc, dst);
    });
    rewriter.replaceOp(op, alloc.getResult());
    return success();
  }
};

/// mg.broadcast (static shapes only): result[i, j] = input[i or 0, j or 0], using 0 for each size-1 input dimension.

Here is the real lowered code for row_sum of a 2x3 matrix. The first nest (over a 2x1 output) fills with zero; the second walks the 2x3 input and adds each element into the cell of its row:

module {
  func.func @f(%arg0: memref<2x3xf64>) -> memref<2x1xf64> {
    %alloc = memref.alloc() : memref<2x1xf64>
    %cst = arith.constant 0.000000e+00 : f64
    affine.for %arg1 = 0 to 2 {
      affine.for %arg2 = 0 to 1 {
        affine.store %cst, %alloc[%arg1, %arg2] : memref<2x1xf64>
      }
    }
    affine.for %arg1 = 0 to 2 {
      affine.for %arg2 = 0 to 3 {
        %c0 = arith.constant 0 : index
        %0 = affine.load %alloc[%arg1, %c0] : memref<2x1xf64>
        %1 = affine.load %arg0[%arg1, %arg2] : memref<2x3xf64>
        %2 = arith.addf %0, %1 : f64
        affine.store %2, %alloc[%arg1, %c0] : memref<2x1xf64>
      }
    }
    return %alloc : memref<2x1xf64>
  }
}

broadcast walks the result and reads the input with index 0 in any dimension where the input has size 1. No arithmetic, only a different way of indexing:

struct BroadcastOpLowering : public OpConversionPattern<mg::BroadcastOp> {
  using OpConversionPattern::OpConversionPattern;
  LogicalResult matchAndRewrite(mg::BroadcastOp op, OpAdaptor adaptor,
                                ConversionPatternRewriter &rewriter) const override {
    auto loc = op.getLoc();
    auto inType = llvm::cast<RankedTensorType>(op.getInput().getType());
    auto outType = llvm::cast<RankedTensorType>(op.getResult().getType());
    auto memrefType = convertTensorToMemref(outType);
    Value in = adaptor.getInput();
    auto alloc = rewriter.create<memref::AllocOp>(loc, memrefType);
    mg::dyn::nest2D(rewriter, loc, outType, {Value(), Value()}, [&](OpBuilder &b, Location l, ValueRange ivs) {
      Value zeroIdx = b.create<arith::ConstantIndexOp>(l, 0);
      Value i = inType.getDimSize(0) == 1 ? zeroIdx : ivs[0];
      Value j = inType.getDimSize(1) == 1 ? zeroIdx : ivs[1];
      Value x = b.create<affine::AffineLoadOp>(l, in, ValueRange{i, j});
      b.create<affine::AffineStoreOp>(l, x, alloc, ivs);
    });
    rewriter.replaceOp(op, alloc.getResult());
    return success();
  }
};

module {
  func.func @f(%arg0: memref<1x3xf64>) -> memref<2x3xf64> {
    %alloc = memref.alloc() : memref<2x3xf64>
    affine.for %arg1 = 0 to 2 {
      affine.for %arg2 = 0 to 3 {
        %c0 = arith.constant 0 : index
        %0 = affine.load %arg0[%c0, %arg2] : memref<1x3xf64>
        affine.store %0, %alloc[%arg1, %arg2] : memref<2x3xf64>
      }
    }
    return %alloc : memref<2x3xf64>
  }
}

The row %arg0[%c0, %arg2] is read with the first index fixed at 0 for every row %arg1 of the result: the single input row is repeated.

The complete change against Chapter 21's files:

The full diff for this chapter (click to expand)

--- ../../part21/code/MgOps.td  2026-10-02 03:06:25.081013492 +0000
+++ MgOps.td    2026-10-02 11:17:33.276282865 +0000
@@ -82,6 +82,34 @@
   let hasVerifier = 1;
 }

+def ReluOp : Mg_Op<"relu", [Pure]> {
+  let summary = "elementwise max(x, 0), the rectified linear unit";
+  let arguments = (ins F64Tensor:$input);
+  let results = (outs F64Tensor:$result);
+  let assemblyFormat = "$input attr-dict `:` type($input) `->` type($result)";
+  let hasVerifier = 1;
+}
+
+def ReduceOp : Mg_Op<"reduce", [Pure]> {
+  let summary = "reduce a matrix along one axis with sum or max";
+  let description = [{
+    `axis = 1` combines the elements of each row, giving an M x 1 result; `axis = 0` combines the
+    elements of each column, giving a 1 x N result. `kind` is "sum" or "max".
+  }];
+  let arguments = (ins F64Tensor:$input, I64Attr:$axis, StrAttr:$kind);
+  let results = (outs F64Tensor:$result);
+  let assemblyFormat = "$input attr-dict `:` type($input) `->` type($result)";
+  let hasVerifier = 1;
+}
+
+def BroadcastOp : Mg_Op<"broadcast", [Pure]> {
+  let summary = "repeat a matrix along every dimension of size 1 to fill a larger static shape";
+  let arguments = (ins F64Tensor:$input);
+  let results = (outs F64Tensor:$result);
+  let assemblyFormat = "$input attr-dict `:` type($input) `->` type($result)";
+  let hasVerifier = 1;
+}
+
 def NegOp : Mg_Op<"neg", [Pure]> {
   let summary = "elementwise negation";
   let arguments = (ins F64Tensor:$input);
--- ../../part21/code/MgDialect.cpp 2026-10-02 03:06:25.081495534 +0000
+++ MgDialect.cpp   2026-10-02 11:17:33.276538753 +0000
@@ -107,6 +107,37 @@
 }
 mlir::LogicalResult NegOp::verify() { return verifyElementwise(*this, {getInput().getType(), getResult().getType()}, "mg.neg"); }

+mlir::LogicalResult ReluOp::verify() { return verifyElementwise(*this, {getInput().getType(), getResult().getType()}, "mg.relu"); }
+
+mlir::LogicalResult ReduceOp::verify() {
+  auto in = llvm::dyn_cast<mlir::RankedTensorType>(getInput().getType());
+  auto out = llvm::dyn_cast<mlir::RankedTensorType>(getResult().getType());
+  if (!in || !out || in.getRank() != 2 || out.getRank() != 2)
+    return emitOpError("mg.reduce only supports rank-2 ranked tensors");
+  if (getAxis() != 0 && getAxis() != 1)
+    return emitOpError("mg.reduce: axis must be 0 or 1, got ") << getAxis();
+  llvm::StringRef k = getKind();
+  if (k != "sum" && k != "max")
+    return emitOpError("mg.reduce: kind must be sum or max, got '") << k << "'";
+  int64_t kept = 1 - getAxis();   // the dimension that survives
+  if (!mg::dyn::compatibleDim(out.getDimSize(kept), in.getDimSize(kept)) || out.getDimSize(getAxis()) != 1)
+    return emitOpError("mg.reduce: reducing axis ") << getAxis() << " of " << in << " must give a result with that axis of size 1, got " << out;
+  return mlir::success();
+}
+
+mlir::LogicalResult BroadcastOp::verify() {
+  auto in = llvm::dyn_cast<mlir::RankedTensorType>(getInput().getType());
+  auto out = llvm::dyn_cast<mlir::RankedTensorType>(getResult().getType());
+  if (!in || !out || in.getRank() != 2 || out.getRank() != 2)
+    return emitOpError("mg.broadcast only supports rank-2 ranked tensors");
+  if (!in.hasStaticShape() || !out.hasStaticShape())
+    return emitOpError("mg.broadcast needs static shapes (dynamic broadcasting is not supported)");
+  for (int d = 0; d < 2; ++d)
+    if (in.getDimSize(d) != out.getDimSize(d) && in.getDimSize(d) != 1)
+      return emitOpError("mg.broadcast: dimension ") << d << " of " << in << " can only grow from size 1, not to " << out;
+  return mlir::success();
+}
+
 mlir::LogicalResult MatmulOp::verify() {
   auto l = llvm::dyn_cast<mlir::RankedTensorType>(getLhs().getType());
   auto r = llvm::dyn_cast<mlir::RankedTensorType>(getRhs().getType());
--- ../../part21/code/LowerToAffine.cpp 2026-10-02 03:12:27.030941608 +0000
+++ LowerToAffine.cpp   2026-10-02 11:17:33.277001200 +0000
@@ -1,3 +1,4 @@
+#include <limits>
 #include "mg/MgDialect.h"
 #include "mg/MgOps.h"
 #include "mg/DynamicShapes.h"
@@ -170,6 +171,91 @@
   }
 };

+/// Chapter 22: relu is a one-input elementwise op, like neg.
+struct ReluOpLowering : public OpConversionPattern<mg::ReluOp> {
+  using OpConversionPattern::OpConversionPattern;
+  LogicalResult matchAndRewrite(mg::ReluOp op, OpAdaptor adaptor,
+                                ConversionPatternRewriter &rewriter) const override {
+    auto loc = op.getLoc();
+    auto tensorType = llvm::cast<RankedTensorType>(op.getResult().getType());
+    auto memrefType = convertTensorToMemref(tensorType);
+    Value in = adaptor.getInput();
+    llvm::SmallVector<Value, 2> extents;
+    for (int64_t d = 0; d < 2; ++d)
+      extents.push_back(tensorType.isDynamicDim(d) ? mg::dyn::extentOf(rewriter, loc, in, d) : Value());
+    auto alloc = mg::dyn::allocFor(rewriter, loc, memrefType, extents);
+    Value zero = rewriter.create<arith::ConstantOp>(loc, rewriter.getF64FloatAttr(0.0));
+    mg::dyn::nest2D(rewriter, loc, tensorType, extents, [&](OpBuilder &b, Location l, ValueRange ivs) {
+      Value x = b.create<affine::AffineLoadOp>(l, in, ivs);
+      b.create<affine::AffineStoreOp>(l, b.create<arith::MaximumFOp>(l, x, zero), alloc, ivs);
+    });
+    rewriter.replaceOp(op, alloc.getResult());
+    return success();
+  }
+};
+
+/// mg.reduce: initialise the result (0 for sum, -infinity for max), then one nest over the INPUT that
+/// folds each element into result[i, 0] (axis 1) or result[0, j] (axis 0).
+struct ReduceOpLowering : public OpConversionPattern<mg::ReduceOp> {
+  using OpConversionPattern::OpConversionPattern;
+  LogicalResult matchAndRewrite(mg::ReduceOp op, OpAdaptor adaptor,
+                                ConversionPatternRewriter &rewriter) const override {
+    auto loc = op.getLoc();
+    auto inType = llvm::cast<RankedTensorType>(op.getInput().getType());
+    auto outType = llvm::cast<RankedTensorType>(op.getResult().getType());
+    auto memrefType = convertTensorToMemref(outType);
+    Value in = adaptor.getInput();
+    int64_t axis = op.getAxis();
+    bool isMax = op.getKind() == "max";
+    llvm::SmallVector<Value, 2> outExtents;
+    for (int64_t d = 0; d < 2; ++d)
+      outExtents.push_back(outType.isDynamicDim(d) ? mg::dyn::extentOf(rewriter, loc, in, d) : Value());
+    auto alloc = mg::dyn::allocFor(rewriter, loc, memrefType, outExtents);
+    double initVal = isMax ? -std::numeric_limits<double>::infinity() : 0.0;
+    Value init = rewriter.create<arith::ConstantOp>(loc, rewriter.getF64FloatAttr(initVal));
+    mg::dyn::nest2D(rewriter, loc, outType, outExtents, [&](OpBuilder &b, Location l, ValueRange ivs) {
+      b.create<affine::AffineStoreOp>(l, init, alloc, ivs);
+    });
+    llvm::SmallVector<Value, 2> inExtents;
+    for (int64_t d = 0; d < 2; ++d)
+      inExtents.push_back(inType.isDynamicDim(d) ? mg::dyn::extentOf(rewriter, loc, in, d) : Value());
+    mg::dyn::nest2D(rewriter, loc, inType, inExtents, [&](OpBuilder &b, Location l, ValueRange ivs) {
+      Value zeroIdx = b.create<arith::ConstantIndexOp>(l, 0);
+      Value i = ivs[0], j = ivs[1];
+      llvm::SmallVector<Value, 2> dst = axis == 1 ? llvm::SmallVector<Value, 2>{i, zeroIdx} : llvm::SmallVector<Value, 2>{zeroIdx, j};
+      Value acc = b.create<affine::AffineLoadOp>(l, alloc, dst);
+      Value x = b.create<affine::AffineLoadOp>(l, in, ivs);
+      Value r = isMax ? Value(b.create<arith::MaximumFOp>(l, acc, x)) : Value(b.create<arith::AddFOp>(l, acc, x));
+      b.create<affine::AffineStoreOp>(l, r, alloc, dst);
+    });
+    rewriter.replaceOp(op, alloc.getResult());
+    return success();
+  }
+};
+
+/// mg.broadcast (static shapes only): result[i, j] = input[i or 0, j or 0], using 0 for each size-1 input dimension.
+struct BroadcastOpLowering : public OpConversionPattern<mg::BroadcastOp> {
+  using OpConversionPattern::OpConversionPattern;
+  LogicalResult matchAndRewrite(mg::BroadcastOp op, OpAdaptor adaptor,
+                                ConversionPatternRewriter &rewriter) const override {
+    auto loc = op.getLoc();
+    auto inType = llvm::cast<RankedTensorType>(op.getInput().getType());
+    auto outType = llvm::cast<RankedTensorType>(op.getResult().getType());
+    auto memrefType = convertTensorToMemref(outType);
+    Value in = adaptor.getInput();
+    auto alloc = rewriter.create<memref::AllocOp>(loc, memrefType);
+    mg::dyn::nest2D(rewriter, loc, outType, {Value(), Value()}, [&](OpBuilder &b, Location l, ValueRange ivs) {
+      Value zeroIdx = b.create<arith::ConstantIndexOp>(l, 0);
+      Value i = inType.getDimSize(0) == 1 ? zeroIdx : ivs[0];
+      Value j = inType.getDimSize(1) == 1 ? zeroIdx : ivs[1];
+      Value x = b.create<affine::AffineLoadOp>(l, in, ValueRange{i, j});
+      b.create<affine::AffineStoreOp>(l, x, alloc, ivs);
+    });
+    rewriter.replaceOp(op, alloc.getResult());
+    return success();
+  }
+};
+
 struct NegOpLowering : public OpConversionPattern<mg::NegOp> {
   using OpConversionPattern::OpConversionPattern;
   LogicalResult matchAndRewrite(mg::NegOp op, OpAdaptor adaptor,
@@ -400,7 +486,8 @@
                  ElementwiseBinaryLowering<mg::SubOp, arith::SubFOp>,
                  ElementwiseBinaryLowering<mg::MulOp, arith::MulFOp>,
                  ElementwiseBinaryLowering<mg::DivOp, arith::DivFOp>,
-                 ScalarOpLowering, NegOpLowering, MatmulOpLowering>(typeConverter, &getContext());
+                 ScalarOpLowering, NegOpLowering, MatmulOpLowering,
+                 ReluOpLowering, ReduceOpLowering, BroadcastOpLowering>(typeConverter, &getContext());
     populateFunctionOpInterfaceTypeConversionPattern<func::FuncOp>(patterns,
                                                                    typeConverter);
     populateCallOpTypeConversionPattern(patterns, typeConverter);
--- ../../part21/code/mgfront.py    2026-10-02 03:17:06.351799178 +0000
+++ mgfront.py  2026-10-02 11:22:48.028568462 +0000
@@ -1,6 +1,8 @@
 #!/usr/bin/env python3
 """mgfront: the Mountain Goat surface-syntax front end. Reads .mg source, writes MLIR (mg dialect) to stdout.

+Built-in functions (Chapter 22): relu(m); row_sum, row_max, row_mean (one value per row, an Mx1 result) and col_sum, col_max, col_mean
+(one value per column, 1xN). Elementwise operators broadcast a static dimension of size 1 (e.g. a 2x3 plus a 1x3).
 Operators (Chapter 21): + - * / are elementwise on matrices ('*' is the Hadamard product) and also combine a matrix
 with a scalar; '@' is the matrix product; unary '-' negates. Scalars are numbers known at compile time.
 Grammar (one statement per line, '#' starts a comment):
@@ -103,12 +105,27 @@
             r = self.fn.tmp()
             self.fn.emit(f'{r} = mg.scalar {m} {{op = "{self.OPNAME[op]}", value = {mlir_float(str(c))} : f64, reversed = {"true" if rev else "false"}}} : {fmt(shp)} -> {fmt(shp)}')
             return r, shp
-        for a, b in zip(s, s2):
-            if a is not None and b is not None and a != b:
-                raise MgError(self.line, f"cannot {self.VERB[op]} shapes {'x'.join(map(str,s))} and {'x'.join(map(str,s2))}")
+        v, s, w, s2 = self.broadcast_pair(op, v, s, w, s2)
         shape = tuple(a if a is not None else b for a, b in zip(s, s2)); r = self.fn.tmp()
         self.fn.emit(f"{r} = mg.{self.OPNAME[op]} {v}, {w} : {fmt(s)}, {fmt(s2)} -> {fmt(shape)}")
         return r, shape
+    def broadcast_pair(self, op, v, s, w, s2):
+        """Make two shapes agree for an elementwise operator. A static dimension of size 1 grows to match the other
+        operand's (explicit mg.broadcast); equal sizes and '?' are left alone; anything else is an error."""
+        target = []
+        for a, b in zip(s, s2):
+            if a is None or b is None or a == b: target.append(b if a is None else a)
+            elif a == 1: target.append(b)
+            elif b == 1: target.append(a)
+            else: raise MgError(self.line, f"cannot {self.VERB[op]} shapes {'x'.join('?' if d is None else str(d) for d in s)} and {'x'.join('?' if d is None else str(d) for d in s2)}: "
+                                           f"dimensions {a} and {b} differ and neither is 1")
+        target = tuple(target)
+        return (*self.grow(v, s, target), *self.grow(w, s2, target))
+    def grow(self, v, s, target):
+        if s == target or all(a is None or a == t for a, t in zip(s, target)): return v, s
+        if None in s or None in target: raise MgError(self.line, "broadcasting needs static shapes; this operand has a '?' dimension")
+        r = self.fn.tmp(); self.fn.emit(f"{r} = mg.broadcast {v} : {fmt(s)} -> {fmt(target)}")
+        return r, target
     def term(self):
         k, v = self.peek()
         if k is None: raise MgError(self.line, "expected an expression, found the end of the line")
@@ -123,11 +140,30 @@
                 shape = (s[1], s[0]); r = self.fn.tmp()
                 self.fn.emit(f"{r} = mg.transpose {x} : {fmt(s)} to {fmt(shape)}")
                 return r, shape
+            if v in self.BUILTINS and self.peek()[1] == "(":
+                return self.builtin(v)
             if self.peek()[1] == "(":
                 return self.call(v)
             if v not in self.fn.env: raise MgError(self.line, f"unknown name {v!r}")
             return self.fn.env[v]
         raise MgError(self.line, f"unexpected {v!r}")
+    BUILTINS = {"relu", "row_sum", "col_sum", "row_max", "col_max", "row_mean", "col_mean"}
+    def builtin(self, name):
+        self.take("("); v, s = self.expr(); self.take(")")
+        if v is None: raise MgError(self.line, f"{name} needs a matrix, not a scalar")
+        if name == "relu":
+            r = self.fn.tmp(); self.fn.emit(f"{r} = mg.relu {v} : {fmt(s)} -> {fmt(s)}"); return r, s
+        axis = 1 if name.startswith("row_") else 0              # row_*: one value per row (axis 1); col_*: one per column (axis 0)
+        kind = "sum" if name.endswith("sum") else "max" if name.endswith("max") else "mean"
+        shape = (s[0], 1) if axis == 1 else (1, s[1])
+        r = self.fn.tmp()
+        self.fn.emit(f'{r} = mg.reduce {v} {{axis = {axis} : i64, kind = "{"sum" if kind == "mean" else kind}"}} : {fmt(s)} -> {fmt(shape)}')
+        if kind == "mean":                                       # mean = sum / count; the count must be known now
+            if s[axis] is None: raise MgError(self.line, f"{name} needs a static size along the averaged axis, not '?'")
+            q = self.fn.tmp()
+            self.fn.emit(f'{q} = mg.scalar {r} {{op = "div", value = {mlir_float(str(s[axis]))} : f64, reversed = false}} : {fmt(shape)} -> {fmt(shape)}')
+            r = q
+        return r, shape
     def call(self, name):
         if name not in self.fn.em.funcs: raise MgError(self.line, f"unknown function {name!r}")
         params, result = self.fn.em.funcs[name]
--- ../../part21/code/mgc   2026-10-02 03:17:06.353500579 +0000
+++ mgc 2026-10-02 11:22:48.030686829 +0000
@@ -5,10 +5,10 @@
 #   mgc lib   prog.mg [-o dir]   compile the 'def's to prog.o plus a C++ header prog.h (no main), for linking into your own program
 #   mgc mlir  prog.mg            print the mg-dialect MLIR the front end produces
 #   mgc ptx   prog.mg            compile to PTX (no launch: this machine has no GPU) and print the PTX
-# Env: MG_OPT (default Chapter 21 build), MGC_KEEP=dir (keep every intermediate file there).
+# Env: MG_OPT (default Chapter 22 build), MGC_KEEP=dir (keep every intermediate file there).
 set -e
 HERE=$(cd "$(dirname "$0")" && pwd)
-MG_OPT=${MG_OPT:-$HERE/../../part21/code/build/mg-opt}
+MG_OPT=${MG_OPT:-$HERE/../../part22/code/build/mg-opt}
 LIBDIR=${MLIR_LIBDIR:-/usr/lib/llvm-18/lib}
 usage() { echo "usage: mgc run|build|lib|mlir|ptx prog.mg [-o exe]" >&2; exit 2; }
 [ $# -ge 2 ] || usage
@@ -20,7 +20,7 @@
 LIBFLAG=; [ "$cmd" = lib ] && LIBFLAG=--lib
 python3 "$HERE/mgfront.py" $LIBFLAG "$src" > "$W/$base.mlir"          # 1. surface syntax -> mg dialect (exit 1 + message on a syntax/shape error)
 [ "$cmd" = mlir ] && { cat "$W/$base.mlir"; exit 0; }
-[ -x "$MG_OPT" ] || { echo "mgc: no mg-opt at $MG_OPT (run part21/code/build.sh)" >&2; exit 2; }
+[ -x "$MG_OPT" ] || { echo "mgc: no mg-opt at $MG_OPT (run part22/code/build.sh)" >&2; exit 2; }
 case $cmd in
   build|run|lib)
     LOWER='--lower-affine --convert-scf-to-cf --mg-lower-assert-to-stderr --convert-arith-to-llvm --finalize-memref-to-llvm --convert-func-to-llvm --reconcile-unrealized-casts'

To rebuild mg-opt:

#!/bin/sh
# Rebuild mg-opt for Chapter 22: Chapter 21's build plus mg.relu, mg.reduce and mg.broadcast (this directory's MgOps.td, MgDialect.cpp, LowerToAffine.cpp override the earlier ones). Produces ./build/mg-opt.
# Needs mlir-18-tools, libmlir-18-dev, llvm-18-dev, cmake, make.
set -e
HERE=$(cd "$(dirname "$0")" && pwd); D=$HERE/../..; T=$HERE/tree
mkdir -p $T/include/mg $T/lib $T/tools
cp $HERE/CMakeLists.txt $T/
cp $D/part3/code/MgDialect.td $HERE/MgOps.td $D/part3/code/MgDialect.h $D/part3/code/MgOps.h $T/include/mg/
cp $D/part14/code/DynamicShapes.h $T/include/mg/
cp $HERE/MgDialect.cpp $T/lib/
cp $D/part14/code/MgBufferizableOpInterfaceImpl.cpp $T/lib/
cp $D/part19/code/AssertToStderr.cpp $T/lib/
cp $HERE/LowerToAffine.cpp $T/lib/
cp $HERE/mg-opt.cpp $T/tools/
mkdir -p $HERE/build/include/mg && cd $HERE/build
cmake $T -DMLIR_DIR=/usr/lib/llvm-18/lib/cmake/mlir -DLLVM_DIR=/usr/lib/llvm-18/lib/cmake/llvm \
  -DCMAKE_BUILD_TYPE=Release -DLLVM_ENABLE_ASSERTIONS=OFF > cmake_out.txt 2>&1
make -j8 2>&1 | tail -2

The front end

Two additions to mgfront.py. First, the broadcast rule, applied to every elementwise operator. For each pair of dimensions it picks the result size (equal sizes, or ?, stay as they are; a static 1 grows to the other side's size; anything else is an error) and then grow inserts an mg.broadcast for any operand that needs one:

    def broadcast_pair(self, op, v, s, w, s2):
        """Make two shapes agree for an elementwise operator. A static dimension of size 1 grows to match the other
        operand's (explicit mg.broadcast); equal sizes and '?' are left alone; anything else is an error."""
        target = []
        for a, b in zip(s, s2):
            if a is None or b is None or a == b: target.append(b if a is None else a)
            elif a == 1: target.append(b)
            elif b == 1: target.append(a)
            else: raise MgError(self.line, f"cannot {self.VERB[op]} shapes {'x'.join('?' if d is None else str(d) for d in s)} and {'x'.join('?' if d is None else str(d) for d in s2)}: "
                                           f"dimensions {a} and {b} differ and neither is 1")
        target = tuple(target)
        return (*self.grow(v, s, target), *self.grow(w, s2, target))
    def grow(self, v, s, target):
        if s == target or all(a is None or a == t for a, t in zip(s, target)): return v, s
        if None in s or None in target: raise MgError(self.line, "broadcasting needs static shapes; this operand has a '?' dimension")
        r = self.fn.tmp(); self.fn.emit(f"{r} = mg.broadcast {v} : {fmt(s)} -> {fmt(target)}")
        return r, target
    def term(self):
        k, v = self.peek()
        if k is None: raise MgError(self.line, "expected an expression, found the end of the line")

Second, the built-in functions (relu, and row_ / col_ followed by sum, max or mean), which turn into mg.relu or mg.reduce, with mean followed by a divide by the static count:

    BUILTINS = {"relu", "row_sum", "col_sum", "row_max", "col_max", "row_mean", "col_mean"}
    def builtin(self, name):
        self.take("("); v, s = self.expr(); self.take(")")
        if v is None: raise MgError(self.line, f"{name} needs a matrix, not a scalar")
        if name == "relu":
            r = self.fn.tmp(); self.fn.emit(f"{r} = mg.relu {v} : {fmt(s)} -> {fmt(s)}"); return r, s
        axis = 1 if name.startswith("row_") else 0              # row_*: one value per row (axis 1); col_*: one per column (axis 0)
        kind = "sum" if name.endswith("sum") else "max" if name.endswith("max") else "mean"
        shape = (s[0], 1) if axis == 1 else (1, s[1])
        r = self.fn.tmp()
        self.fn.emit(f'{r} = mg.reduce {v} {{axis = {axis} : i64, kind = "{"sum" if kind == "mean" else kind}"}} : {fmt(s)} -> {fmt(shape)}')
        if kind == "mean":                                       # mean = sum / count; the count must be known now
            if s[axis] is None: raise MgError(self.line, f"{name} needs a static size along the averaged axis, not '?'")
            q = self.fn.tmp()
            self.fn.emit(f'{q} = mg.scalar {r} {{op = "div", value = {mlir_float(str(s[axis]))} : f64, reversed = false}} : {fmt(shape)} -> {fmt(shape)}')
            r = q
        return r, shape

And here is what the front end emits for a whole layer, relu(x @ w + bias), where bias is a 1x2 row and x @ w is 2x2. The mg.broadcast on the third line was inserted automatically:

func.func @layer(%x: tensor<2x3xf64>, %w: tensor<3x2xf64>, %bias: tensor<1x2xf64>) -> tensor<2x2xf64> {
  %t1 = mg.matmul %x, %w : tensor<2x3xf64>, tensor<3x2xf64> -> tensor<2x2xf64>
  %t2 = mg.broadcast %bias : tensor<1x2xf64> -> tensor<2x2xf64>
  %t3 = mg.add %t1, %t2 : tensor<2x2xf64>, tensor<2x2xf64> -> tensor<2x2xf64>
  %t4 = mg.relu %t3 : tensor<2x2xf64> -> tensor<2x2xf64>
  func.return %t4 : tensor<2x2xf64>

Thirteen examples

Each block shows the program, its real output, and the exit status. Memref base addresses (0x…) are masked.

1. relu

################ ./mgc run examples/01_relu.mg
  | # relu(x) = max(x, 0): negative numbers become 0, the rest are unchanged.
  | print relu([[1, -2, 3], [-4, 5, -6]])
  ----- output
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [2, 3] strides = [3, 1] data = 
[[1,   0,   3], 
 [0,   5,   0]]
  ----- exit status 0

2. Row sums and column sums

################ ./mgc run examples/02_row_and_col_sums.mg
  | # row_sum: one value per ROW (an Mx1 result). col_sum: one value per COLUMN (a 1xN result).
  | let a = [[1, 2, 3], [4, 5, 6]]
  | print row_sum(a)
  | print col_sum(a)
  ----- output
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [2, 1] strides = [1, 1] data = 
[[6], 
 [15]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 3] strides = [3, 1] data = 
[[5,   7,   9]]
  ----- exit status 0

row_sum gives a 2x1 column (one total per row: 1+2+3 = 6, 4+5+6 = 15); col_sum gives a 1x3 row (one total per column: 5, 7, 9).

3. Max and mean

################ ./mgc run examples/03_max_and_mean.mg
  | # The same pair for max and for mean. A mean is a sum divided by the count, so the count must be known.
  | let a = [[1, 9, 3], [8, 5, 6]]
  | print row_max(a)
  | print col_max(a)
  | print row_mean(a)
  | print col_mean(a)
  ----- output
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [2, 1] strides = [1, 1] data = 
[[9], 
 [8]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 3] strides = [3, 1] data = 
[[8,   9,   6]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [2, 1] strides = [1, 1] data = 
[[4.33333], 
 [6.33333]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 3] strides = [3, 1] data = 
[[4.5,   7,   4.5]]
  ----- exit status 0

row_mean of the first row [1, 9, 3] is 13/3 = 4.33333; the printer shows six significant digits.

4. Broadcasting a bias

################ ./mgc run examples/04_broadcast_bias.mg
  | # Broadcasting: a size-1 dimension is repeated to fit. Add a 1x3 row to every row of a 2x3 matrix.
  | let a = [[1, 2, 3], [4, 5, 6]]
  | let bias = [[100, 200, 300]]
  | print a + bias
  | let col = [[10], [20]]
  | print a + col                     # a 2x1 column is repeated across the columns
  | print a * [[2]]                   # a 1x1 matrix is repeated in both directions
  ----- output
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [2, 3] strides = [3, 1] data = 
[[101,   202,   303], 
 [104,   205,   306]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [2, 3] strides = [3, 1] data = 
[[11,   12,   13], 
 [24,   25,   26]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [2, 3] strides = [3, 1] data = 
[[2,   4,   6], 
 [8,   10,   12]]
  ----- exit status 0

A 1x3 row is added to every row; a 2x1 column is added to every column; a 1x1 matrix scales everything ([[2]] behaves like the scalar 2, which is why a scalar is simpler when it is all you need).

5. Centering columns

################ ./mgc run examples/05_center_columns.mg
  | # Centering: subtract each column's mean, so every column of the result averages to zero.
  | let a = [[1, 10], [2, 20], [3, 60]]
  | let centered = a - col_mean(a)
  | print centered
  | print col_sum(centered)
  ----- output
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [3, 2] strides = [2, 1] data = 
[[-1,   -20], 
 [0,   -10], 
 [1,   30]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 2] strides = [2, 1] data = 
[[0,   0]]
  ----- exit status 0

a - col_mean(a) subtracts each column's mean (2 and 30 here) from that column. The check after it, col_sum(centered), is exactly [[0, 0]]: a centered column adds up to zero. This is a property that holds whatever the lowering does, so it is a good independent test of reductions, broadcasting and mean together.

6. Normalizing rows

################ ./mgc run examples/06_normalize_rows.mg
  | # Divide each row by its own sum: every row of the result adds up to 1 (a probability-like distribution).
  | let a = [[1, 1, 2], [5, 0, 5]]
  | let p = a / row_sum(a)
  | print p
  | print row_sum(p)
  ----- output
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [2, 3] strides = [3, 1] data = 
[[0.25,   0.25,   0.5], 
 [0.5,   0,   0.5]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [2, 1] strides = [1, 1] data = 
[[1], 
 [1]]
  ----- exit status 0

Dividing each row by its own sum makes every row add to 1 (the second output confirms it).

7. A neural-network layer

################ ./mgc run examples/07_mlp_layer.mg
  | # One neural-network layer: relu(x @ w + bias). Two samples, three inputs, two outputs.
  | def layer(x: tensor[2x3], w: tensor[3x2], bias: tensor[1x2]) = relu(x @ w + bias)
  | print layer([[1, 2, 3], [-1, -2, -3]], [[1, -1], [0, 1], [1, 0]], [[0.5, -10]])
  ----- output
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [2, 2] strides = [2, 1] data = 
[[4.5,   0], 
 [0,   0]]
  ----- exit status 0

relu(x @ w + bias): a matrix product, a broadcast add and an activation, in one expression. Check the top-left value by hand: row [1, 2, 3] times column [1, 0, 1] is 4, plus bias 0.5 is 4.5, and relu keeps it. The top-right: -1 + 2 + 0 = 1, plus -10 is -9, and relu clamps it to 0. The second sample's whole row is negative, so it becomes zeros.

8. Two layers

################ ./mgc run examples/08_two_layers.mg
  | # Two layers, and the largest output of each row, the "prediction" of a tiny network.
  | def net(x: tensor[2x2]) = row_max(relu(relu(x @ [[1, -1], [2, 1]] + [[0, 1]]) @ [[1], [-1]] + [[0]]) @ [[1, 2]])
  | print net([[1, 1], [-1, 2]])
  ----- output
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [2, 1] strides = [1, 1] data = 
[[4], 
 [0]]
  ----- exit status 0

Two layers and a final row_max, the largest value in each row. By hand: [1, 1] goes through layer one to [3, 1], layer two to [2] and finally to [2, 4], whose max is 4; the second row ends as [0, 0], max 0.

9. Reductions on dynamic shapes

################ ./mgc run examples/09_dynamic_reductions.mg
  | # Reductions work on dynamic shapes: one compiled function, any number of rows or columns.
  | def totals(a: tensor[?x?]) = row_sum(a)
  | def peaks(a: tensor[?x?]) = col_max(a)
  | print totals([[1, 2, 3], [4, 5, 6]])
  | print totals([[7], [8], [9], [10]])
  | print peaks([[1, 9], [5, 2], [3, 3]])
  | print relu(-[[1, 2], [3, 4]] + 2)
  ----- output
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [2, 1] strides = [1, 1] data = 
[[6], 
 [15]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [4, 1] strides = [1, 1] data = 
[[7], 
 [8], 
 [9], 
 [10]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 2] strides = [2, 1] data = 
[[5,   9]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [2, 2] strides = [2, 1] data = 
[[1,   0], 
 [0,   0]]
  ----- exit status 0

totals and peaks take tensor[?x?], so each is compiled once and handles 2x3, 4x1 and 3x2 inputs. The last line is relu(-[[1,2],[3,4]] + 2): unary minus, then a scalar add, then relu.

10. ? is never broadcast

################ ./mgc run examples/11_dynamic_no_broadcast.mg
  | # A '?' dimension is never broadcast: the sizes must be equal when the program runs.
  | # The first call has one row, which matches the 1x3 bias. The second has two rows, so the run-time check aborts.
  | def plus_bias(a: tensor[?x3]) = a + [[1, 2, 3]]
  | print plus_bias([[10, 20, 30]])
  | print plus_bias([[10, 20, 30], [40, 50, 60]])
  ----- output
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 3] strides = [3, 1] data = 
[[11,   22,   33]]
mg.add: operand shapes differ at runtime in dimension 0
Aborted
  ----- exit status 134

The first call has one row, which matches the 1x3 bias exactly, so it works. The second call has two rows, and because that dimension is ? the program compares sizes at run time instead of broadcasting, and aborts with the message on stderr (Chapters 14 and 19). If you want the broadcast, the row count must be a static number in the type.

11. Broadcasting from the left

################ ./mgc run examples/12_broadcast_left.mg
  | # Broadcasting works whichever side the small operand is on, and for non-commutative operators the order is kept.
  | let a = [[1, 2, 3], [4, 5, 6]]
  | let row = [[10, 20, 30]]
  | print row + a                     # the 1x3 row is on the LEFT
  | print row - a                     # 10-1, 20-2, ...: left minus right, not the reverse
  | print [[100], [200]] - a          # a 2x1 column on the left
  | print [[2]] * a
  ----- output
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [2, 3] strides = [3, 1] data = 
[[11,   22,   33], 
 [14,   25,   36]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [2, 3] strides = [3, 1] data = 
[[9,   18,   27], 
 [6,   15,   24]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [2, 3] strides = [3, 1] data = 
[[99,   98,   97], 
 [196,   195,   194]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [2, 3] strides = [3, 1] data = 
[[2,   4,   6], 
 [8,   10,   12]]
  ----- exit status 0

The small operand can be on either side, and operand order is preserved: row - a is 10-1, 20-2, 30-3 (left minus right), not the reverse. This example was added after a mutation test found it missing (see Tests below).

12. The max of negative numbers

################ ./mgc run examples/13_negative_max.mg
  | # The largest of a set of negative numbers is a negative number, not 0. (A max that started from 0 would get this wrong.)
  | let a = [[-3, -1], [-5, -2]]
  | print row_max(a)
  | print col_max(a)
  ----- output
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [2, 1] strides = [1, 1] data = 
[[-1], 
 [-2]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 2] strides = [2, 1] data = 
[[-3,   -1]]
  ----- exit status 0

The largest of -3 and -1 is -1, not 0. This looks too obvious to test, and for exactly that reason it was not tested at first: a max reduction that started from 0 instead of negative infinity gets every earlier example right (they all have positive numbers) and only fails here. This example was added after a mutation test found it missing (see Tests below).

13. The GPU path

################ ./mgc ptx examples/10_gpu_layer.mg
  | # The GPU path (compile-only) for a layer with broadcast, reduction and relu. Use: mgc ptx 10_gpu_layer.mg
  | def f(x: tensor[2x3], bias: tensor[1x3]) = row_sum(relu(x + bias))
  ----- kernel entry points and the arithmetic (one kernel per loop nest)
.visible .entry f_kernel(
ld.global.f64   %fd1, [%rd11];
st.global.f64   [%rd14], %fd1;
.visible .entry f_kernel(
ld.global.f64   %fd1, [%rd16];
ld.global.f64   %fd2, [%rd17];
add.rn.f64  %fd3, %fd1, %fd2;
st.global.f64   [%rd18], %fd3;
.visible .entry f_kernel(
ld.global.f64   %fd2, [%rd14];
max.NaN.f64     %fd3, %fd2, %fd1;
st.global.f64   [%rd15], %fd3;
.visible .entry f_kernel(
st.global.f64   [%rd12], %fd1;
.visible .entry f_kernel(
ld.global.f64   %fd1, [%rd11];
ld.global.f64   %fd2, [%rd13];
add.rn.f64  %fd3, %fd1, %fd2;
st.global.f64   [%rd11], %fd3;
ld.global.f64   %fd4, [%rd13+8];
add.rn.f64  %fd5, %fd3, %fd4;
st.global.f64   [%rd11], %fd5;
ld.global.f64   %fd6, [%rd13+16];
add.rn.f64  %fd7, %fd5, %fd6;
st.global.f64   [%rd11], %fd7;

These are real PTX instructions from mgc ptx (compile-only: nothing was launched or assembled with ptxas). Five loop nests make five kernels: broadcast (a load and a store), add, relu (note max.NaN.f64, the NaN-propagating maximum), the reduction's zero-fill, and the reduction's accumulate. In the last, the three add.rn.f64 instructions are the sequential loop over the three columns, unrolled (the +8 and +16 offsets are the second and third columns): each GPU thread owns one row and sums its three elements. The compiler parallelized the row loop and correctly left the reduction loop sequential.

Errors

################ ./mgc run examples/errors/bad_broadcast.mg
  | # 2x3 and 3x2: the dimensions differ and neither is 1, so they cannot be broadcast.
  | print [[1, 2, 3], [4, 5, 6]] + [[1, 2], [3, 4], [5, 6]]
  ----- output
errors/bad_broadcast.mg:2: error: cannot add shapes 2x3 and 3x2: dimensions 2 and 3 differ and neither is 1
  ----- exit status 1

################ ./mgc run examples/errors/dynamic_mean.mg
  | # A mean divides by a count, and the count of rows is not known when the size is '?'.
  | def f(a: tensor[?x?]) = row_mean(a)
  ----- output
errors/dynamic_mean.mg:2: error: row_mean needs a static size along the averaged axis, not '?'
  ----- exit status 1

Using it from C++

# Chapter 22 operations as functions callable from C++.
def relu_m(a: tensor[?x?]) = relu(a)
def row_sums(a: tensor[?x?]) = row_sum(a)
def col_maxes(a: tensor[?x?]) = col_max(a)
def layer(x: tensor[4x6], w: tensor[6x3], bias: tensor[1x3]) = relu(x @ w + bias)
def center(a: tensor[5x4]) = a - col_mean(a)

A C++ program calls each function on matrices larger than any example (7x5, 4x6 by 6x3, 5x4) and compares every result with a plain C++ loop it wrote independently:

// C++ calling the Chapter 22 functions, compared against plain C++ loops on matrices bigger than any example.
#include <algorithm>
#include <cmath>
#include <cstdio>
#include "ops.h"

static mg::Matrix filled(int64_t r, int64_t c, int seed) {
  std::vector<double> d(r * c);
  for (size_t i = 0; i < d.size(); i++) d[i] = (double)(((i + seed) * 37) % 23) - 11.0;
  return mg::Matrix(r, c, d);
}
static bool same(const mg::Matrix &x, const mg::Matrix &y) {
  if (x.rows != y.rows || x.cols != y.cols) return false;
  for (size_t i = 0; i < x.data.size(); i++) if (std::fabs(x.data[i] - y.data[i]) > 1e-9) return false;
  return true;
}
static void report(const char *what, bool ok) { std::printf("%-44s %s\n", what, ok ? "matches" : "DIFFERS"); }

int main() {
  mg::Matrix a = filled(7, 5, 1);
  mg::Matrix r(7, 5, std::vector<double>(35));
  for (size_t i = 0; i < 35; i++) r.data[i] = std::max(a.data[i], 0.0);
  report("relu_m 7x5", same(mg::relu_m(a), r));

  mg::Matrix rs(7, 1, std::vector<double>(7, 0.0));
  for (int64_t i = 0; i < 7; i++) for (int64_t j = 0; j < 5; j++) rs(i, 0) += a(i, j);
  report("row_sums 7x5", same(mg::row_sums(a), rs));

  mg::Matrix cm(1, 5, std::vector<double>(5, -1e300));
  for (int64_t i = 0; i < 7; i++) for (int64_t j = 0; j < 5; j++) cm(0, j) = std::max(cm(0, j), a(i, j));
  report("col_maxes 7x5", same(mg::col_maxes(a), cm));

  mg::Matrix x = filled(4, 6, 2), w = filled(6, 3, 3), b = filled(1, 3, 4), want(4, 3, std::vector<double>(12, 0.0));
  for (int64_t i = 0; i < 4; i++) for (int64_t j = 0; j < 3; j++) {
    double acc = 0; for (int64_t k = 0; k < 6; k++) acc += x(i, k) * w(k, j);
    want(i, j) = std::max(acc + b(0, j), 0.0);
  }
  report("layer: relu(x @ w + bias), 4x6 by 6x3", same(mg::layer(x, w, b), want));

  mg::Matrix m = filled(5, 4, 5), c(5, 4, std::vector<double>(20));
  for (int64_t j = 0; j < 4; j++) {
    double mean = 0; for (int64_t i = 0; i < 5; i++) mean += m(i, j); mean /= 5;
    for (int64_t i = 0; i < 5; i++) c(i, j) = m(i, j) - mean;
  }
  report("center: subtract each column's mean, 5x4", same(mg::center(m), c));
  return 0;
}

#!/bin/sh
# Compile ops.mg with mgc, then build and run the C++ program against it. Output: run_out.txt
HERE=$(cd "$(dirname "$0")" && pwd); W=$HERE/../work/cpp; mkdir -p $W
$HERE/../mgc lib $HERE/ops.mg -o $W
cp $W/ops.h $HERE/ops_generated.h
clang++-18 -std=c++17 -Wall -I$W $HERE/app.cpp $W/ops.o -o $W/app
echo "--- ./app"; $W/app

wrote work/cpp/ops.o and work/cpp/ops.h
--- ./app
relu_m 7x5                                   matches
row_sums 7x5                                 matches
col_maxes 7x5                                matches
layer: relu(x @ w + bias), 4x6 by 6x3        matches
center: subtract each column's mean, 5x4     matches

All five match.

Tests

Ten new test files in test/broadcast/ (the suite now has 88 tests). As in earlier chapters, program tests run the real example files.

Test	What it checks
verifier-reduce	bad axis, bad kind, and a result that does not collapse the axis: exact messages
verifier-broadcast	a dimension that cannot grow, and a dynamic shape: exact messages
roundtrip-and-lowering	the three ops parse and print; the reduction has a -infinity start, two nests and no run-time check; the broadcast reads index 0
values	relu, sums, max, mean, broadcasting (including from the left): exact numbers
identities	centered columns sum to 0; normalized rows sum to 1
networks	the layer and the two-layer network
dynamic-reductions	reductions on ? shapes; ? is never broadcast (abort with the message)
front-end-errors	the two front-end errors, exact text
ptx-layer	five kernels, max.NaN.f64, a sequential reduction
cpp-interop-reductions	the C++ program builds with -Wall -Werror and all five comparisons match

prove_tests_can_fail.sh runs them against Chapter 21's mg-opt (none of the new ops exist) and five broken copies of the front end:

#!/bin/sh
# Show that the Chapter 22 tests (test/broadcast) can fail. Output: prove_tests_can_fail_out.txt
#  (1) against Chapter 21's mg-opt, which has none of the new ops;
#  (2) against broken copies of the front end (restored on exit);
#  (3) lowering bugs are in lowering_mutation.sh (rebuilds mg-opt for each).
# READ THIS FIRST: this script runs the Chapter 22 tests against an older build and against deliberately broken copies of the front end.
#   FAIL lines (and non-zero exit statuses) in this script's output are the EXPECTED result: each broken copy SHOULD make at least one test fail (a FAIL line means a test caught the bug). The "baseline" sections must show all tests passing.
#   The unmodified, current build is the baseline and must show no failures. A mutation or old build that makes NO test fail would be the problem.
echo "NOTE: this script deliberately runs broken or older code. FAIL lines below are EXPECTED: they show the tests can detect the problem. The baseline (unmodified current build) must show none."
HERE=$(cd "$(dirname "$0")" && pwd); LIT=$HERE/../../part15/code; W=$HERE/work/prove; rm -rf $W; mkdir -p $W
run() { label=$1; mgopt=$2
  printf '%s\n' "### $label"
  MG_OPT=$mgopt MG_TEST_TMP=$W/out sh $LIT/run_lit.sh --filter ':: broadcast/' 2>&1 | grep -E "^(FAIL)|  Passed|  Failed" | sed 's#IR Forge / Mountain Goat :: ##'
}
run "baseline: Chapter 22 build" $HERE/build/mg-opt
run "Chapter 21 build (none of the new ops exist)" $HERE/../../part21/code/build/mg-opt
cp $HERE/mgfront.py $W/mgfront.orig; cp $HERE/mgc $W/mgc.orig
restore() { cp $W/mgfront.orig $HERE/mgfront.py; cp $W/mgc.orig $HERE/mgc; }
trap restore EXIT
mutate() {  # label file sed-expr
  restore; cp $2 $W/before; sed -i "$3" $2
  if cmp -s $2 $W/before; then echo "### $1: MUTATION DID NOT APPLY"; return; fi
  run "front-end mutation: $1" $HERE/build/mg-opt
}
mutate "row_* and col_* use the opposite axis"          $HERE/mgfront.py 's/axis = 1 if name.startswith("row_") else 0/axis = 0 if name.startswith("row_") else 1/'
mutate "mean divides by the wrong dimension's count"    $HERE/mgfront.py 's/mlir_float(str(s\[axis\]))/mlir_float(str(s[1 - axis]))/'
mutate "relu is emitted as negation"                    $HERE/mgfront.py 's/mg.relu {v} : {fmt(s)} -> {fmt(s)}/mg.neg {v} : {fmt(s)} -> {fmt(s)}/'
mutate "a size-1 left operand is not broadcast"         $HERE/mgfront.py 's/            elif a == 1: target.append(b)/            elif False: target.append(b)/'
mutate "a 1x1 right operand is never broadcast"         $HERE/mgfront.py 's/            elif b == 1: target.append(a)/            elif False: target.append(a)/'
restore; echo "### restored"; run "baseline again" $HERE/build/mg-opt

NOTE: this script deliberately runs broken or older code. FAIL lines below are EXPECTED: they show the tests can detect the problem. The baseline (unmodified current build) must show none.
### baseline: Chapter 22 build
  Passed  : 10 (11.11%)
### Chapter 21 build (none of the new ops exist)
FAIL: broadcast/dynamic-reductions.mlir (1 of 10)
FAIL: broadcast/identities.mlir (2 of 10)
FAIL: broadcast/cpp-interop-reductions.mlir (3 of 10)
FAIL: broadcast/networks.mlir (4 of 10)
FAIL: broadcast/roundtrip-and-lowering.mlir (5 of 10)
FAIL: broadcast/ptx-layer.mlir (6 of 10)
FAIL: broadcast/values.mlir (7 of 10)
FAIL: broadcast/verifier-reduce.mlir (8 of 10)
FAIL: broadcast/verifier-broadcast.mlir (10 of 10)
  Passed  :  1 (1.11%)
  Failed  :  9 (10.00%)
### front-end mutation: row_* and col_* use the opposite axis
FAIL: broadcast/identities.mlir (1 of 10)
FAIL: broadcast/values.mlir (2 of 10)
FAIL: broadcast/networks.mlir (3 of 10)
FAIL: broadcast/dynamic-reductions.mlir (4 of 10)
FAIL: broadcast/ptx-layer.mlir (9 of 10)
FAIL: broadcast/cpp-interop-reductions.mlir (10 of 10)
  Passed  :  4 (4.44%)
  Failed  :  6 (6.67%)
### front-end mutation: mean divides by the wrong dimension's count
FAIL: broadcast/values.mlir (4 of 10)
FAIL: broadcast/identities.mlir (6 of 10)
FAIL: broadcast/cpp-interop-reductions.mlir (10 of 10)
  Passed  :  7 (7.78%)
  Failed  :  3 (3.33%)
### front-end mutation: relu is emitted as negation
FAIL: broadcast/values.mlir (1 of 10)
FAIL: broadcast/dynamic-reductions.mlir (2 of 10)
FAIL: broadcast/ptx-layer.mlir (3 of 10)
FAIL: broadcast/networks.mlir (4 of 10)
FAIL: broadcast/cpp-interop-reductions.mlir (10 of 10)
  Passed  :  5 (5.56%)
  Failed  :  5 (5.56%)
### front-end mutation: a size-1 left operand is not broadcast
FAIL: broadcast/values.mlir (10 of 10)
  Passed  :  9 (10.00%)
  Failed  :  1 (1.11%)
### front-end mutation: a 1x1 right operand is never broadcast
FAIL: broadcast/networks.mlir (1 of 10)
FAIL: broadcast/cpp-interop-reductions.mlir (2 of 10)
FAIL: broadcast/identities.mlir (3 of 10)
FAIL: broadcast/ptx-layer.mlir (4 of 10)
FAIL: broadcast/values.mlir (10 of 10)
  Passed  :  5 (5.56%)
  Failed  :  5 (5.56%)
### restored
### baseline again
  Passed  : 10 (11.11%)

Two runs found holes in the tests. The first time the front-end mutations ran, the mutation "a size-1 left operand is not broadcast" made no test fail (Passed: 10, no failures): every example put the small operand on the right (a + bias), so the left-hand branch of the front end's broadcast rule was never exercised. The left-broadcast example (12_broadcast_left.mg, section 11 below) and a new check in values were added, the mutation was rerun, and it is now caught. A test suite is evidence only for what it exercises; a mutation is the way to find what it does not.

Lowering bugs need mg-opt rebuilt, so lowering_mutation.sh injects each into a copy of the source tree:

#!/bin/sh
# Rebuild mg-opt with one bug injected into LowerToAffine.cpp and run the ops tests against it.
# Each build is a full rebuild in work/mutN (about two minutes). Output: lowering_mutation_out.txt
# READ THIS FIRST: this script rebuilds mg-opt with one deliberate bug injected into the lowering at a time (a full rebuild each) and runs the tests.
#   FAIL lines (and non-zero exit statuses) in this script's output are the EXPECTED result: each bug SHOULD make at least one test fail (a FAIL line means a test caught it). A section with NO FAIL line would be a test gap.
#   The unmodified, current build is the baseline and must show no failures. A mutation or old build that makes NO test fail would be the problem.
echo "NOTE: this script deliberately runs broken or older code. FAIL lines below are EXPECTED: they show the tests can detect the problem. The baseline (unmodified current build) must show none."
HERE=$(cd "$(dirname "$0")" && pwd); LIT=$HERE/../../part15/code; M=$HERE/work/mut; rm -rf $M; mkdir -p $M
n=0
try() {  # label sed-expr
  n=$((n+1)); d=$M/$n; mkdir -p $d/build; cp -r $HERE/tree $d/tree; rm -rf $d/tree/build
  cp $d/tree/lib/LowerToAffine.cpp $d/before.cpp; sed -i "$2" $d/tree/lib/LowerToAffine.cpp
  if cmp -s $d/before.cpp $d/tree/lib/LowerToAffine.cpp; then echo "### $1: MUTATION DID NOT APPLY"; return; fi
  ( cd $d/build && cmake $d/tree -DMLIR_DIR=/usr/lib/llvm-18/lib/cmake/mlir -DLLVM_DIR=/usr/lib/llvm-18/lib/cmake/llvm -DCMAKE_BUILD_TYPE=Release -DLLVM_ENABLE_ASSERTIONS=OFF >/dev/null 2>&1 && mkdir -p include/mg && make -j8 >/dev/null 2>&1 )
  [ -x $d/build/mg-opt ] || { echo "### $1: BUILD FAILED"; return; }
  printf '%s\n' "### lowering mutation: $1"
  MG_OPT=$d/build/mg-opt MG_TEST_TMP=$d/out sh $LIT/run_lit.sh --filter ':: broadcast/' 2>&1 | grep -E "^(FAIL)|  Passed|  Failed" | sed 's#IR Forge / Mountain Goat :: ##'
}
try "max reduction starts at 0 instead of -infinity"      's/double initVal = isMax ? -std::numeric_limits<double>::infinity() : 0.0;/double initVal = 0.0;/'
try "relu does not clamp (max of x with itself)"          's/b.create<arith::MaximumFOp>(l, x, zero)/b.create<arith::MaximumFOp>(l, x, x)/'
try "broadcast reads index 0 of the wrong dimension"      's/Value j = inType.getDimSize(1) == 1 ? zeroIdx : ivs\[1\];/Value j = inType.getDimSize(0) == 1 ? zeroIdx : ivs[1];/'
try "reduce over axis 1 writes into the wrong cell"        's/axis == 1 ? llvm::SmallVector<Value, 2>{i, zeroIdx}/axis == 1 ? llvm::SmallVector<Value, 2>{zeroIdx, i}/'

### lowering mutation: max reduction starts at 0 instead of -infinity
FAIL: broadcast/roundtrip-and-lowering.mlir (4 of 10)
FAIL: broadcast/cpp-interop-reductions.mlir (9 of 10)
FAIL: broadcast/values.mlir (10 of 10)
  Passed  :  7 (7.95%)
  Failed  :  3 (3.41%)
### lowering mutation: relu does not clamp (max of x with itself)
FAIL: broadcast/dynamic-reductions.mlir (2 of 10)
FAIL: broadcast/networks.mlir (3 of 10)
FAIL: broadcast/ptx-layer.mlir (5 of 10)
FAIL: broadcast/values.mlir (9 of 10)
FAIL: broadcast/cpp-interop-reductions.mlir (10 of 10)
  Passed  :  5 (5.68%)
  Failed  :  5 (5.68%)
### lowering mutation: broadcast reads index 0 of the wrong dimension
FAIL: broadcast/identities.mlir (2 of 10)
FAIL: broadcast/networks.mlir (3 of 10)
FAIL: broadcast/roundtrip-and-lowering.mlir (4 of 10)
FAIL: broadcast/cpp-interop-reductions.mlir (9 of 10)
FAIL: broadcast/values.mlir (10 of 10)
  Passed  :  5 (5.68%)
  Failed  :  5 (5.68%)
### lowering mutation: reduce over axis 1 writes into the wrong cell
FAIL: broadcast/roundtrip-and-lowering.mlir (5 of 10)
  Passed  :  9 (10.23%)
  Failed  :  1 (1.14%)

The second hole. The first lowering mutation, "max reduction starts at 0 instead of -infinity", was run before the negative-max example (13_negative_max.mg, section 12 below) existed. It was caught by only two of the ten test files (the lowering-structure test, which looks for the -infinity constant, and the C++ comparison, whose generated data happens to include a column of all-negative numbers). Every hand-sized example passed, because none had a column whose maximum is negative. The first run's result:

### lowering mutation: max reduction starts at 0 instead of -infinity
FAIL: broadcast/roundtrip-and-lowering.mlir (5 of 10)
FAIL: broadcast/cpp-interop-reductions.mlir (9 of 10)
  Passed  :  8 (9.09%)
  Failed  :  2 (2.27%)

The negative-max example and a check in values were added, and the whole mutation run was repeated; the listing above is the repeat, in which that mutation now also fails values.

The third: a mutation that changes nothing you can see. "Reduce over axis 1 writes into the wrong cell" stores each row's running total at result[0, i] instead of result[i, 0]. It passed every test, and no test of the numbers could ever catch it. The reason is that the result of a row reduction is M by 1, and in row-major memory the element at [0, i] has the same address as [i, 0] (the offset is row * 1 + column, and one of the two indices is 0). The program writes outside the declared shape (in the abstract language that is out of bounds) but into exactly the memory that the correct version would have used, so the outputs are identical. A bug whose effect is invisible in the output is called an equivalent mutant, and behavioral tests cannot see it. A structural check can: roundtrip-and-lowering now also requires that the row reduction loads and stores %alloc[%arg, %c0] (row index first, constant zero second), and with that check the mutation fails one test. It is a fair question whether such a check tests behavior or implementation. Here it does the latter, deliberately, because the behavior is indistinguishable and the implementation is wrong: on a different layout, or after a later optimization that trusts the declared bounds, the bug would surface.

The full suite against the Chapter 22 build:

Testing Time: 4.28s

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

Limits and what is not established

• Rank 2, f64 only, as before.
• Broadcasting is static only. A ? dimension is never broadcast; it must match at run time. No broadcasting for @.
• Mean needs a static count; row_mean on a ?-sized axis is rejected.
• Only sum and max reductions (plus mean, built from sum). No min, argmax, product, or reduction over both axes at once.
• relu and NaN: relu(NaN) is NaN by design. No test feeds a NaN in (the language has no NaN literal; you can produce one with 0/0, but this chapter does not test it).
• Sum order and precision: a reduction adds in loop order. Floating-point addition is not associative, so other orders (a GPU, a library) can differ in the last digits. The C++ comparison tolerates 1e-9.
• No performance claim. The reductions are plain loops, and the GPU kernels are one thread per output row with a sequential inner loop; none of this was timed.
• GPU: PTX produced and inspected, never assembled, loaded or launched. Correctness on a GPU is not established.
• Bufferization (Chapter 6): not extended; the new operations lower through the hand-written conversion that mgc uses.
• Platform: x86-64 Linux, LLVM 18.1.3, POSIX only.

Reproducing

cd docs/part22/code
./build.sh && ./demo.sh > demo_out.txt && ./show.sh && cpp/run.sh
./prove_tests_can_fail.sh            # front-end mutations (seconds)
./lowering_mutation.sh               # lowering mutations (a full rebuild each: roughly 10 minutes)
cd ../../part15/code && ./run_lit.sh # 88 tests

Chapter summary

• Three new operations: mg.relu, mg.reduce (sum or max along an axis) and mg.broadcast; the front end adds relu, row_ and col_ sum, max and mean, and broadcasts a static size-1 dimension automatically.
• A reduction needs two loop nests (initialize, then fold the input into the output) and shrinks one dimension to size 1, which is what makes a - col_mean(a) possible.
• Broadcasting changes shapes, not arithmetic: the front end inserts an explicit mg.broadcast, and later stages see ordinary equal-shape operations. A ? dimension is never broadcast; it is checked at run time.
• relu(x @ w + bias) runs: a real neural-network layer, and a two-layer network, with values checked by hand.
• C++ calls it and matches independent plain loops; PTX is produced for all of it.
• A mutation test found a real gap (broadcasting from the left was untested); it is fixed and covered.

Self-check questions

Each answer is collapsed; try the question first.

1. What does broadcasting do, and what is the rule for when two shapes can be broadcast together?

   Answer

   A dimension of size 1 is conceptually repeated to match the other operand, so [[10, 20, 30]] (1x3) can be added to a 2x3 matrix as if it were two stacked copies. Per dimension, the sizes must be equal or one of them must be 1; otherwise the shapes cannot be broadcast. It changes only shapes, not the arithmetic. In Mountain Goat the rule applies only to static dimensions.

2. Why is ? never broadcast?

   Answer

   A ? dimension's size is unknown when the program is compiled, so the compiler cannot know whether it is 1 (and so repeatable) or not. Guessing would silently change the meaning of a program depending on its inputs. Instead the dimension must equal the other operand's at run time, and a mismatch aborts with a message (example 10). To get broadcasting, write a static size in the type.

3. What is the difference between row_sum(a) and col_sum(a), and what shapes do they return for a 2x3 matrix?

   Answer

   row_sum gives one total per row, a 2x1 result (the three columns collapse). col_sum gives one total per column, a 1x3 result (the two rows collapse). The name says what you get one of; the reduced dimension becomes size 1.

4. Why does mg.reduce lower to two loop nests instead of one?

   Answer

   Each output element combines several input elements, so it must start from a known value and then accumulate. The first nest stores the starting value (0 for a sum, negative infinity for a max) into every output cell; the second walks the input and folds each element into its output cell. A single nest that stored each output once would not have anywhere to accumulate.

5. Why does a max reduction start at negative infinity rather than at 0?

   Answer

   max(x, -infinity) = x for every x, so the starting value never affects the answer. Starting at 0 would be wrong whenever every element is negative: the result would be 0, a number that is not in the data at all. (This is one of the lowering mutations; see the results above for whether the tests catch it.)

6. What is a - col_mean(a) for a 3x2 matrix a, shape by shape?

   Answer

   col_mean(a) is 1x2 (one mean per column). Subtracting a 1x2 from a 3x2 broadcasts the 1x2 across the three rows, so every element has its own column's mean subtracted. The result is 3x2, and each column of it sums to zero.

7. Why is row_mean an error on a tensor[?x?] but row_sum is not?

   Answer

   A mean is a sum divided by a count, and the front end bakes the count into the program as a constant. With a ? size the count does not exist until run time, so there is no constant to divide by. A sum needs no count, so it works for any size.

8. The first mutation run left "a size-1 left operand is not broadcast" uncaught. What does that tell you about the tests, and what was done?

   Answer

   It showed that no test exercised the branch of the broadcast rule where the left operand is the one that grows: every example wrote the small operand on the right, so the branch could be deleted with no test noticing. An example with the small operand on the left (including a non-commutative operator, so order matters) and a matching check were added, and the same mutation now fails a test.

9. In the PTX, the reduction's accumulate kernel has three add.rn.f64 instructions in a row for a 2x3 input. What are they?

   Answer

   They are the three iterations of the inner loop over the columns, unrolled because the column count (3) is a constant. Each GPU thread handles one row and adds that row's three elements into the output cell in order. The row loop was made parallel (one thread per row); the reduction loop was left sequential, because parallelizing a loop that accumulates into one cell would be a race.

10. relu is lowered with a NaN-propagating maximum. What is the observable difference from a maximum that ignores NaN?

    Answer

    With a NaN input, a propagating maximum returns NaN, so the problem stays visible in the output, while a NaN-ignoring maximum would return 0 and hide it. This chapter chose propagation (and the GPU's max.NaN). No test feeds a NaN to relu, so that behavior is by design and by the instruction's definition, not something the tests demonstrate.