diff --git a/src/jacobi/core.py b/src/jacobi/core.py index f6be06b..5e6ead8 100644 --- a/src/jacobi/core.py +++ b/src/jacobi/core.py @@ -266,8 +266,10 @@ def propagate( Returns ------- y, ycov - y is the result of fn(x) + y is the result of fn(x). ycov is the propagated covariance matrix. + If ycov is a matrix, unless y is a number. In that case, ycov is also + reduced to a number. """ x = np.array(x) y = fn(x) @@ -291,4 +293,8 @@ def propagate( raise ValueError("x and cov have incompatible shapes") ycov = np.einsum("il,kl,l" if xcov_nd == 1 else "ij,kl,jl", jac, jac, xcov) + + if np.ndim(y) == 0: + ycov = ycov[0, 0] + return y, ycov diff --git a/test/test_propagate.py b/test/test_propagate.py index 1724b45..3544b6b 100644 --- a/test/test_propagate.py +++ b/test/test_propagate.py @@ -44,6 +44,8 @@ def fn(x): assert_allclose(y, fn(x)) jac = np.ones((1, len(x))) assert_allclose(ycov, np.linalg.multi_dot([jac, xcov, jac.T])) + assert np.ndim(y) == 0 + assert np.ndim(ycov) == 0 def test_11():