From ffdd80e1e6e9f29cb37b9bd667a1b2e83a3ce83d Mon Sep 17 00:00:00 2001 From: Tomoki Uda Date: Mon, 23 Mar 2026 22:25:57 +0900 Subject: [PATCH 1/6] Fix coef_from_cov_grad: validate inputs and handle singular covariances 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 --- src/ellphi/grad.py | 30 ++++++++++++++++++++++++++++-- 1 file changed, 28 insertions(+), 2 deletions(-) diff --git a/src/ellphi/grad.py b/src/ellphi/grad.py index 35ffcc9..bc12a4e 100644 --- a/src/ellphi/grad.py +++ b/src/ellphi/grad.py @@ -183,11 +183,37 @@ def coef_from_cov_grad( if cov_arr.ndim == 2: cov_arr = cov_arr[np.newaxis, :, :] + 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]: + raise ValueError("Covariance matrices must be square") + n, d = centers.shape + + if cov_arr.shape[-1] != d: + raise ValueError("Centre dimensionality and covariance size must agree") s2 = scale**2 - # Forward pass - inv_cov = np.linalg.inv(cov_arr) # (n, d, d) + # Forward pass – handle singular covariances the same way coef_from_cov + # does: return NaN coefficients so callers can detect degeneracy. + try: + inv_cov = np.linalg.inv(cov_arr) # (n, d, d) + except np.linalg.LinAlgError: + m = (d + 1) * (d + 2) // 2 + coefs = np.full((n, m), np.nan, dtype=float) + + def vjp_nan( + grad_coefs: np.ndarray, + ) -> tuple[np.ndarray, np.ndarray]: + return ( + np.zeros_like(centers), + np.zeros_like(cov_arr), + ) + + return coefs, vjp_nan + A = inv_cov / s2 # (n, d, d) b = -np.einsum("nij,nj->ni", A, centers) # (n, d) c = np.einsum("ni,nij,nj->n", centers, A, centers) # (n,) From 679d15421565570ee97225d94f1cac162b983d94 Mon Sep 17 00:00:00 2001 From: Tomoki Uda Date: Mon, 23 Mar 2026 22:45:27 +0900 Subject: [PATCH 2/6] Fix black formatting in grad.py Co-Authored-By: Claude Opus 4.6 --- src/ellphi/grad.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/src/ellphi/grad.py b/src/ellphi/grad.py index bc12a4e..0d56092 100644 --- a/src/ellphi/grad.py +++ b/src/ellphi/grad.py @@ -184,9 +184,7 @@ def coef_from_cov_grad( cov_arr = cov_arr[np.newaxis, :, :] if centers.shape[0] != cov_arr.shape[0]: - raise ValueError( - "Mismatch between number of centres and covariance matrices" - ) + raise ValueError("Mismatch between number of centres and covariance matrices") if cov_arr.shape[-1] != cov_arr.shape[-2]: raise ValueError("Covariance matrices must be square") From 38f914b46dd7223816b8bddd6f47f31488db2cc3 Mon Sep 17 00:00:00 2001 From: Tomoki Uda Date: Mon, 23 Mar 2026 22:49:09 +0900 Subject: [PATCH 3/6] Add tests for coef_from_cov_grad input validation and singular fallback 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 --- tests/test_coef_from_cov_grad.py | 39 ++++++++++++++++++++++++++++++++ 1 file changed, 39 insertions(+) diff --git a/tests/test_coef_from_cov_grad.py b/tests/test_coef_from_cov_grad.py index ff419c2..0e01d44 100644 --- a/tests/test_coef_from_cov_grad.py +++ b/tests/test_coef_from_cov_grad.py @@ -185,3 +185,42 @@ def test_grad_shapes(rng): grad_X, grad_cov = vjp(g) assert grad_X.shape == (n, d) assert grad_cov.shape == (n, d, d) + + +# --------------------------------------------------------------------------- +# Input validation +# --------------------------------------------------------------------------- + +import pytest + + +def test_batch_size_mismatch(): + """Mismatched centre/covariance batch sizes raise ValueError.""" + with pytest.raises(ValueError, match="Mismatch"): + coef_from_cov_grad(np.zeros((3, 2)), np.eye(2)[np.newaxis]) + + +def test_non_square_cov(): + """Non-square covariance raises ValueError.""" + with pytest.raises(ValueError, match="square"): + coef_from_cov_grad(np.zeros((2, 2)), np.zeros((2, 3, 2))) + + +def test_dimension_mismatch(): + """Centre/covariance dimension mismatch raises ValueError.""" + with pytest.raises(ValueError, match="dimensionality"): + coef_from_cov_grad(np.zeros((2, 2)), np.eye(3)[np.newaxis].repeat(2, axis=0)) + + +def test_singular_cov_returns_nan(): + """Singular covariance returns NaN coefficients, matching coef_from_cov.""" + centers = np.zeros((1, 2)) + cov_sing = np.zeros((1, 2, 2)) + + coefs, vjp = coef_from_cov_grad(centers, cov_sing) + ref = coef_from_cov(centers, cov_sing) + np.testing.assert_array_equal(np.isnan(coefs), np.isnan(ref)) + + grad_X, grad_cov = vjp(np.ones_like(coefs)) + np.testing.assert_array_equal(grad_X, 0.0) + np.testing.assert_array_equal(grad_cov, 0.0) From f567009c4ba0002bec91ce6ce1f899cbb721c5ae Mon Sep 17 00:00:00 2001 From: Tomoki Uda Date: Mon, 23 Mar 2026 22:51:21 +0900 Subject: [PATCH 4/6] Move pytest import to top of file (flake8 E402) Co-Authored-By: Claude Opus 4.6 --- tests/test_coef_from_cov_grad.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/tests/test_coef_from_cov_grad.py b/tests/test_coef_from_cov_grad.py index 0e01d44..fd36314 100644 --- a/tests/test_coef_from_cov_grad.py +++ b/tests/test_coef_from_cov_grad.py @@ -3,6 +3,7 @@ from __future__ import annotations import numpy as np +import pytest from ellphi import coef_from_cov_grad from ellphi.geometry import coef_from_cov @@ -191,9 +192,6 @@ def test_grad_shapes(rng): # Input validation # --------------------------------------------------------------------------- -import pytest - - def test_batch_size_mismatch(): """Mismatched centre/covariance batch sizes raise ValueError.""" with pytest.raises(ValueError, match="Mismatch"): From ff2616d01dc1ff19ca5e63198d2f470604d8675e Mon Sep 17 00:00:00 2001 From: Tomoki Uda Date: Mon, 23 Mar 2026 22:53:17 +0900 Subject: [PATCH 5/6] Fix black formatting in test_coef_from_cov_grad.py Co-Authored-By: Claude Opus 4.6 --- tests/test_coef_from_cov_grad.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/test_coef_from_cov_grad.py b/tests/test_coef_from_cov_grad.py index fd36314..96929e6 100644 --- a/tests/test_coef_from_cov_grad.py +++ b/tests/test_coef_from_cov_grad.py @@ -192,6 +192,7 @@ def test_grad_shapes(rng): # Input validation # --------------------------------------------------------------------------- + def test_batch_size_mismatch(): """Mismatched centre/covariance batch sizes raise ValueError.""" with pytest.raises(ValueError, match="Mismatch"): From 78d2300195fbc385ff443f6f885e796c5aa48655 Mon Sep 17 00:00:00 2001 From: Tomoki Uda Date: Mon, 23 Mar 2026 23:18:11 +0900 Subject: [PATCH 6/6] Return NaN gradients (not zeros) for singular covariances 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 --- src/ellphi/grad.py | 4 ++-- tests/test_coef_from_cov_grad.py | 4 ++-- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/src/ellphi/grad.py b/src/ellphi/grad.py index 0d56092..741e57c 100644 --- a/src/ellphi/grad.py +++ b/src/ellphi/grad.py @@ -206,8 +206,8 @@ def vjp_nan( grad_coefs: np.ndarray, ) -> tuple[np.ndarray, np.ndarray]: return ( - np.zeros_like(centers), - np.zeros_like(cov_arr), + np.full_like(centers, np.nan), + np.full_like(cov_arr, np.nan), ) return coefs, vjp_nan diff --git a/tests/test_coef_from_cov_grad.py b/tests/test_coef_from_cov_grad.py index 96929e6..56bf8d5 100644 --- a/tests/test_coef_from_cov_grad.py +++ b/tests/test_coef_from_cov_grad.py @@ -221,5 +221,5 @@ def test_singular_cov_returns_nan(): np.testing.assert_array_equal(np.isnan(coefs), np.isnan(ref)) grad_X, grad_cov = vjp(np.ones_like(coefs)) - np.testing.assert_array_equal(grad_X, 0.0) - np.testing.assert_array_equal(grad_cov, 0.0) + assert np.all(np.isnan(grad_X)) + assert np.all(np.isnan(grad_cov))