Skip to content

Preserve broadcast operand shapes in solve gradients - #831

Open
AHMETHAKANBEZIR1 wants to merge 1 commit into
HIPS:masterfrom
AHMETHAKANBEZIR1:fix/solve-broadcast-gradients
Open

AHMETHAKANBEZIR1 wants to merge 1 commit into
HIPS:masterfrom
AHMETHAKANBEZIR1:fix/solve-broadcast-gradients

Conversation

@AHMETHAKANBEZIR1

Copy link
Copy Markdown

Differentiating numpy.linalg.solve with broadcast operands can fail or return a gradient with the wrong shape. For example, NumPy 2 evaluates solve(A, b) for A.shape == (3, 2, 2) and b.shape == (2,), but either reverse-mode gradient raises a core-dimension error. A shared matrix with batched matrix right-hand sides also receives an unreduced, incorrectly shaped gradient.

Determine vector right-hand sides using the installed NumPy version's solve semantics, promote them explicitly for the adjoint solve, and reduce each operand's gradient back to its original broadcast shape with the existing unbroadcast_f helper. Constant Python-list right-hand sides continue to work.

Validation:

  • Five new broadcast cases covering both arguments and second-order reverse-mode gradients: unchanged master has five failures; the patch passes all five. Existing solve tests and a constant-list control also pass (11 focused tests total).
  • NumPy 2.5.4 / SciPy 1.18.1 / Python 3.12.14: full suite 651 passed, 14 warnings.
  • Separate NumPy 1.26.4 / SciPy 1.13.1 environment: full linear-algebra module 67 passed, 2 skipped (the new NumPy-2-only vector semantics).
  • All configured pre-commit hooks and git diff --check pass. Separate complex, all-entry finite-difference checks cover both NumPy versions and agree within 4e-11.

AI disclosure: OpenAI Codex autonomously prepared the code, tests, and this description. Validation and an additional AI review were performed; there has been no independent human review.

This branch has not been deployed

No deployments
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.

1 participant