Skip to content

Fix CuPy batched solve shape mismatch in _estimate_transfer_function#76

Merged
edeno merged 2 commits into
Eden-Kramer-Lab:masterfrom
tomHep10:fix/cupy-batched-solve-shape
May 12, 2026
Merged

Fix CuPy batched solve shape mismatch in _estimate_transfer_function#76
edeno merged 2 commits into
Eden-Kramer-Lab:masterfrom
tomHep10:fix/cupy-batched-solve-shape

Conversation

@tomHep10

Copy link
Copy Markdown
Contributor

Fixes CuPy GPU crash in _estimate_transfer_function where a non-batched identity matrix (M, M) was passed to xp.linalg.solve with a batched a of shape (..., 1, M, M).

Fix: broadcast identity to match H_0.shape before solving.

Confirmed working on NVIDIA B200 GPU with LFP dataset.

@emilyps14 emilyps14 requested a review from edeno May 12, 2026 14:10
@emilyps14

Copy link
Copy Markdown
Collaborator

I checked this and it looks reasonable to me. I'm fairly certain that all of the annoying assert changes are just formatting. Want @edeno to doublecheck the CuPy main edits.

Broadcast the identity RHS to match H's batch dimensions before calling xp.linalg.solve in connectivity routines (regularized inverse path). This prevents CuPy batched-solve shape mismatch crashes (affecting DTF/PDC/GPDC/DDTF on GPU). Also: add test test_regularized_solve_rhs_matches_batched_lhs to assert RHS/LHS shapes match, update reshaping tuple style, apply minor test assertion formatting cleanups, and expand the CHANGELOG entry to describe the fix.
@codecov

codecov Bot commented May 12, 2026

Copy link
Copy Markdown

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 88.35%. Comparing base (a5d2983) to head (02eca6d).

Additional details and impacted files
@@            Coverage Diff             @@
##           master      #76      +/-   ##
==========================================
- Coverage   88.65%   88.35%   -0.30%     
==========================================
  Files           8        8              
  Lines        1789     1786       -3     
==========================================
- Hits         1586     1578       -8     
- Misses        203      208       +5     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

@edeno

edeno commented May 12, 2026

Copy link
Copy Markdown
Contributor

Hi @tomHep10,

Thank you for catching this and for your contribution. I hope you don't mind, but I reverted some changes the ruff formatter would have complained about, added a test, and fixed another spot this would be an issue.

@edeno edeno merged commit 24baf8b into Eden-Kramer-Lab:master May 12, 2026
11 of 12 checks passed
@tomHep10

Copy link
Copy Markdown
Contributor Author

@edeno Not at all! Thank you, glad this helped.

edeno added a commit that referenced this pull request May 12, 2026
Patch release fixing CuPy batched-solve shape mismatch in transfer
function and MVAR Fourier coefficient computation (PR #76).
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.

3 participants