Skip to content

linalg: add kron - #262

Open
axiom-of-choice wants to merge 1 commit into
sbryngelson:mainfrom
axiom-of-choice:feat/kron
Open

linalg: add kron#262
axiom-of-choice wants to merge 1 commit into
sbryngelson:mainfrom
axiom-of-choice:feat/kron

Conversation

@axiom-of-choice

Copy link
Copy Markdown
Contributor

Adds kron(A, B) to aneforge.linalg, built as a broadcast-multiply of expanded views plus a reshape (no new kernel):

  • A.expand_dims((1, 3)) * B.expand_dims((0, 2))[m, p, n, q], then reshape to [m*p, n*q].

Closes #126.

Testing

  • test_kron in tests/test_linalg.py (two non-square shapes) compares against np.kron; on-device M2 Pro: passed (max abs err 4.9e-4, ~half ULP fp16).
  • Off-device suite: 407 passed, 2 skipped.
  • ruff, pylint 2-space gate, pyright aneforge, compileall clean.

@sbryngelson sbryngelson left a comment

Copy link
Copy Markdown
Owner

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Verified on device (M5): the kron test passes, and the full linalg suite is unaffected.

The expand_dims/broadcast-multiply/reshape construction is the right shape for this, and
_check_precision=False matches the existing convention in this module (linalg.py:36, 360, 716).
Nothing to change.

One scheduling note: #259, #260 and this all append to the same __all__ block and the same tail
of linalg.py, so whichever lands first will force a rebase on the other two.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

linalg: kron (Kronecker product)

2 participants