Try first

You have the definition. Write matmul(A, B) before reading on. Three nested loops, and the hardest part is deciding which index goes where.

The library

function shape(A) {
  return [A.length, A[0].length];
}

function matmul(A, B) {
  const [ar, ac] = shape(A);
  const [br, bc] = shape(B);
  if (ac !== br) {
    throw new Error("cannot multiply " + ar + "x" + ac + " by " + br + "x" + bc);
  }
  const out = [];
  for (let i = 0; i < ar; i++) {
    const row = new Array(bc).fill(0);
    for (let k = 0; k < bc; k++) {
      let sum = 0;
      for (let j = 0; j < ac; j++) sum += A[i][j] * B[j][k];
      row[k] = sum;
    }
    out.push(row);
  }
  return out;
}

function transpose(A) {
  const [r, c] = shape(A);
  const out = [];
  for (let j = 0; j < c; j++) {
    const row = new Array(r);
    for (let i = 0; i < r; i++) row[i] = A[i][j];
    out.push(row);
  }
  return out;
}

The three loops in matmul map onto the definition directly. i walks the rows of the answer, k walks its columns, and j is the shared dimension that gets summed away and does not appear in the output.

The shape check is the most valuable line in this lesson. Without it, mismatched shapes give undefined in the arithmetic and NaN in the output, and you find out about it four functions later with no clue where it started. Fail loudly and immediately.

A few more

function zeros(r, c) {
  return Array.from({ length: r }, () => new Array(c).fill(0));
}

function colVector(a) {
  return a.map(v => [v]);      // length n array -> n by 1 matrix
}

function flatten(A) {
  return A.map(row => row[0]); // n by 1 matrix -> length n array
}

function addScalar(A, k) {
  return A.map(row => row.map(v => v + k));
}

colVector and flatten convert between a plain array and a one column matrix. That conversion is pure bookkeeping and it exists only because matrix multiplication insists on two dimensions. Keeping it in two small named functions is better than scattering [v] and [0] through the rest of the code.

Test it

const A = [[1, 2], [3, 4]];
const B = [[5, 6], [7, 8]];
log(JSON.stringify(matmul(A, B)));      // [[19,22],[43,50]]
log(JSON.stringify(matmul(B, A)));      // [[23,34],[31,46]]
log(JSON.stringify(transpose(A)));      // [[1,3],[2,4]]

const I = [[1, 0], [0, 1]];
log(JSON.stringify(matmul(A, I)));      // A unchanged

try { matmul(A, [[1, 2, 3]]); } catch (err) { log(err.message); }

Four checks worth keeping. The two orders give different answers, which is the point from the last lesson. The identity matrix leaves things alone, which is a good sanity check that your indices are the right way round. And the bad shape throws instead of returning nonsense.

What to watch

This implementation is the clear one, not the fast one. It is three nested loops with no attention to how the numbers sit in memory. A serious library reorders the loops so that memory is read in long straight runs, splits the work into blocks that fit in cache, and hands the result to hardware built for it. That is worth ten to a hundred times the speed.

That gap is most of the honest answer to why real machine learning does not run in loops like these. Lesson 3115 comes back to this.

Exercises

  1. Add a check to transpose that every row has the same length, and test it with a ragged array.
  2. Confirm that transpose(transpose(A)) gives back A for a non square matrix.
  3. Time matmul on two 200 by 200 random matrices. Then 400 by 400. The time should go up about eightfold, and the exercise is to say why eight.
  4. Swap the j and k loops so that j is outermost. The answer is the same. Time both on 300 by 300 and see whether the order changed anything.
PreviousNext