Main function for computing matmul with Strassen in JAX

jax_strassen.py

def f(a: BlockMatrix, b: BlockMatrix) -> BlockMatrix:
    """Multiplies block matrices `a` and `b`."""
    n = len(a)
    result = [[None] * n for _ in range(n)]
    for alpha in range(rank):
      left = None
      for i in range(n):
        for j in range(n):
          if factors[0][i, j, alpha] != 0:
            curr = factors[0][i, j, alpha] * a[i][j]
            if left is None:
              left = curr
            else:
              left += curr
      right = None
      for j in range(n):
        for k in range(n):
          if factors[1][j, k, alpha] != 0:
            curr = factors[1][j, k, alpha] * b[j][k]
            if right is None:
              right = curr
            else:
              right += curr

      matrix_product = left @ right

      for i in range(n):
        for k in range(n):
          if factors[2][i, k, alpha] != 0:
            curr = factors[2][i, k, alpha] * matrix_product
            if result[i][k] is None:
              result[i][k] = curr
            else:
              result[i][k] += curr
    return result