linalg: add kron

aneforge/linalg.py

@@ -528,6 +528,7 @@ def eigvals(A, iters: int = 60):
   "qr", "cholesky", "lu", "lu_pivoted", "solve", "solve_triangular", "inv", "lstsq", "expm",
   "eigh", "eigvals", "generalized_eigh", "dominant_eig",
   "svd", "dominant_svd", "svdvals_topk", "randomized_svd", "pca",
+  "kron",
 ]
 
 
@@ -719,6 +720,23 @@ def norm(A, order="fro"):
   return v
 
 
+def kron(A, B):
+  """Kronecker product A (x) B by broadcast-multiply of expanded views; result [m*p, n*q]."""
+  A16 = np.asarray(A, f16); B16 = np.asarray(B, f16)
+  if A16.ndim != 2 or B16.ndim != 2:
+    raise ValueError(f"linalg.kron: expected 2-D matrices; got {A16.shape} and {B16.shape}")
+  m, n = A16.shape; p, q = B16.shape
+  At = af.input((m, n)); Bt = af.input((p, q))
+  A4 = At.expand_dims((1, 3))
+  B4 = Bt.expand_dims((0, 2))
+  C = A4 * B4
+  out = C.reshape(m * p, n * q)
+  net = af.compile(out, _check_precision=False)
+  Y = net(A16, B16)
+  net.release()
+  return np.asarray(Y, np.float32)
+
+
 def expm(A, order: int = 8):
   """exp(A) by scaling and squaring: exp(A) = exp(A / 2^s)^(2^s), with the inner exponential a
   Taylor sum of `order` terms.

tests/test_linalg.py

@@ -475,3 +475,20 @@ def test_lstsq_rejects():
   with pytest.raises(ValueError): L.lstsq(np.zeros(4, f16), np.zeros(4, f16))          # not 2-D
   with pytest.raises(ValueError): L.lstsq(np.zeros((10, 4), f16), np.zeros(9, f16))    # b rows != m
   with pytest.raises(np.linalg.LinAlgError): L.lstsq(np.zeros((10, 4), f16), np.zeros(10, f16))
+
+
+def test_kron():
+  r = np.random.default_rng(0)
+  max_err = 0.0
+  for m, n, p, q in [(2, 3, 4, 5), (3, 2, 2, 4)]:
+    A16 = r.standard_normal((m, n)).astype(np.float16)
+    B16 = r.standard_normal((p, q)).astype(np.float16)
+    got = L.kron(A16, B16)
+    ref = np.kron(A16.astype(np.float64), B16.astype(np.float64))
+    assert got.shape == (m * p, n * q)
+    got16 = got.astype(np.float16)
+    ref16 = ref.astype(np.float16)
+    err = np.abs(got16.astype(np.float64) - ref16.astype(np.float64)).max()
+    max_err = max(max_err, err)
+    assert np.allclose(got16, ref16, atol=5e-4, rtol=0), f"kron({m},{n},{p},{q}) max abs err {err}"
+  assert max_err <= 5e-4