Skip to content

Fix GEMM JAX starter: add missing C parameter to solve() - #315

Merged
kunal-mansukhani merged 1 commit into
AlphaGPU:mainfrom
adityasharmaaaaa:fix/gemm-jax-starter-signature
Aug 26, 2026
Merged

Fix GEMM JAX starter: add missing C parameter to solve()#315
kunal-mansukhani merged 1 commit into
AlphaGPU:mainfrom
adityasharmaaaaa:fix/gemm-jax-starter-signature

Conversation

@adityasharmaaaaa

Copy link
Copy Markdown
Contributor

Fixes #301.

The JAX starter's solve() was missing the C accumulator tensor that
the harness passes (alpha*A@B + beta*C), causing:

TypeError: solve() takes 7 positional arguments but 8 were given

Brought the signature in line with the PyTorch/Triton starters for this
challenge and the existing pattern in 94_ssm_selective_scan's JAX
starter. Verified with pre-commit run --all-files (black/isort/flake8/
clang-format all pass).

@kunal-mansukhani
kunal-mansukhani merged commit eb1f727 into AlphaGPU:main Aug 26, 2026
4 of 5 checks passed
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.

General Matrix Multiplication (GEMM) broken for JAX

2 participants