Fix coef_from_cov_grad: validate inputs and handle singular covariances - #102
Conversation
Add batch-size, squareness, and dimensionality checks matching coef_from_cov() so malformed inputs raise ValueError instead of silently broadcasting or returning NaN. Singular but well-formed covariances now return NaN coefficients with zero-gradient VJP, matching the forward-only fallback behavior. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: ffdd80e1e6
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| if centers.shape[0] != cov_arr.shape[0]: | ||
| raise ValueError( | ||
| "Mismatch between number of centres and covariance matrices" | ||
| ) | ||
| if cov_arr.shape[-1] != cov_arr.shape[-2]: |
There was a problem hiding this comment.
Validate input rank before indexing shape tuple
The new validation path indexes centers.shape[0] and cov_arr.shape[-2] without first checking rank, so malformed inputs like scalar X or 1-D cov now raise IndexError instead of the intended ValueError. This makes error handling inconsistent for callers that rely on catching ValueError for bad shapes and contradicts the commit’s goal of robust input validation.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
This is already handled. The rank promotion at lines 181–184 ensures centers is at least 2-D and cov_arr is at least 3-D before the validation runs:
if centers.ndim == 1:
centers = centers[np.newaxis, :]
if cov_arr.ndim == 2:
cov_arr = cov_arr[np.newaxis, :, :]Inputs with ranks outside the documented contract ((n, d) / (d,) for X, (n, d, d) / (d, d) for cov) are not supported — this matches coef_from_cov() in geometry.py:318-324 which uses identical promotion logic and has no additional rank checks.
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
Codecov Report✅ All modified and coverable lines are covered by tests. 📢 Thoughts on this report? Let us know! |
Cover the new validation paths: batch size mismatch, non-square covariance, dimension mismatch (all ValueError), and singular covariance returning NaN with zero-gradient VJP. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
Zero gradients falsely signal a stationary point to optimisers. NaN gradients are consistent with the NaN coefficients and correctly indicate an undefined derivative. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
Add batch-size, squareness, and dimensionality checks matching coef_from_cov() so malformed inputs raise ValueError instead of silently broadcasting or returning NaN. Singular but well-formed covariances now return NaN coefficients with zero-gradient VJP, matching the forward-only fallback behavior.