Merge branch 'pymc-devs:main' into issue-173

Author: cerrussell

pytensor/link/jax/dispatch/nlinalg.py

@@ -3,7 +3,7 @@
 from pytensor.link.jax.dispatch import jax_funcify
 from pytensor.tensor.blas import BatchedDot
 from pytensor.tensor.math import Dot, MaxAndArgmax
-from pytensor.tensor.nlinalg import SVD, Det, Eig, Eigh, MatrixInverse, QRFull
+from pytensor.tensor.nlinalg import SVD, Det, Eig, Eigh, MatrixInverse, QRFull, SLogDet
 
 
 @jax_funcify.register(SVD)
@@ -25,6 +25,14 @@ def det(x):
     return det
 
 
+@jax_funcify.register(SLogDet)
+def jax_funcify_SLogDet(op, **kwargs):
+    def slogdet(x):
+        return jnp.linalg.slogdet(x)
+
+    return slogdet
+
+
 @jax_funcify.register(Eig)
 def jax_funcify_Eig(op, **kwargs):
     def eig(x):

pytensor/link/numba/dispatch/nlinalg.py

@@ -18,6 +18,7 @@
     MatrixInverse,
     MatrixPinv,
     QRFull,
+    SLogDet,
 )
 
 
@@ -58,6 +59,25 @@ def det(x):
     return det
 
 
+@numba_funcify.register(SLogDet)
+def numba_funcify_SLogDet(op, node, **kwargs):
+
+    out_dtype_1 = node.outputs[0].type.numpy_dtype
+    out_dtype_2 = node.outputs[1].type.numpy_dtype
+
+    inputs_cast = int_to_float_fn(node.inputs, out_dtype_1)
+
+    @numba_basic.numba_njit
+    def slogdet(x):
+        sign, det = np.linalg.slogdet(inputs_cast(x))
+        return (
+            numba_basic.direct_cast(sign, out_dtype_1),
+            numba_basic.direct_cast(det, out_dtype_2),
+        )
+
+    return slogdet
+
+
 @numba_funcify.register(Eig)
 def numba_funcify_Eig(op, node, **kwargs):
 

pytensor/tensor/extra_ops.py

@@ -1643,6 +1643,11 @@ def make_node(self, a, *shape):
 
         shape, static_shape = at.infer_static_shape(shape)
 
+        if len(shape) < a.ndim:
+            raise ValueError(
+                f"Broadcast target shape has {len(shape)} dims, which is shorter than input with {a.ndim} dims"
+            )
+
         out = TensorType(dtype=a.type.dtype, shape=static_shape)()
 
         # Attempt to prevent in-place operations on this view-based output
@@ -1686,9 +1691,12 @@ def infer_shape(self, fgraph, node, ins_shapes):
         return [node.inputs[1:]]
 
     def c_code(self, node, name, inputs, outputs, sub):
+        inp_dims = node.inputs[0].ndim
+        out_dims = node.outputs[0].ndim
+        new_dims = out_dims - inp_dims
+
         (x, *shape) = inputs
         (out,) = outputs
-        ndims = len(shape)
         fail = sub["fail"]
 
         # TODO: Could just use `PyArray_Return`, no?
@@ -1701,20 +1709,34 @@ def c_code(self, node, name, inputs, outputs, sub):
 
         src = (
             """
-            npy_intp itershape[%(ndims)s] = {%(dims_array)s};
+            npy_intp itershape[%(out_dims)s] = {%(dims_array)s};
 
+            NpyIter *iter;
             PyArrayObject *ops[1] = {%(x)s};
             npy_uint32 flags = NPY_ITER_MULTI_INDEX | NPY_ITER_REFS_OK | NPY_ITER_ZEROSIZE_OK;
             npy_uint32 op_flags[1] = {NPY_ITER_READONLY};
             PyArray_Descr *op_dtypes[1] = {NULL};
-            int oa_ndim = %(ndims)s;
+            int oa_ndim = %(out_dims)s;
             int* op_axes[1] = {NULL};
             npy_intp buffersize = 0;
 
-            NpyIter *iter = NpyIter_AdvancedNew(
+            for(int i = 0; i < %(inp_dims)s; i++)
+            {
+                if ((PyArray_DIMS(%(x)s)[i] != 1) && (PyArray_DIMS(%(x)s)[i] != itershape[i + %(new_dims)s]))
+                {
+                    PyErr_Format(PyExc_ValueError,
+                                 "Shape mismatch in broadcast_to: target shape[%%i] = %%lld is incompatible with input shape = %%lld.",
+                                 i,
+                                 (long long int) itershape[i + %(new_dims)s],
+                                 (long long int) PyArray_DIMS(%(x)s)[i]
+                    );
+                    %(fail)s
+                }
+            }
+
+            iter = NpyIter_AdvancedNew(
                 1, ops, flags, NPY_CORDER, NPY_NO_CASTING, op_flags, op_dtypes, oa_ndim, op_axes, itershape, buffersize
             );
-
             %(out)s = NpyIter_GetIterView(iter, 0);
 
             if(%(out)s == NULL){
@@ -1733,7 +1755,7 @@ def c_code(self, node, name, inputs, outputs, sub):
         return src
 
     def c_code_cache_version(self):
-        return (1,)
+        return (2,)
 
 
 broadcast_to_ = BroadcastTo()

pytensor/tensor/nlinalg.py

@@ -231,6 +231,39 @@ def __str__(self):
 det = Det()
 
 
+class SLogDet(Op):
+    """
+    Compute the log determinant and its sign of the matrix. Input should be a square matrix.
+    """
+
+    __props__ = ()
+
+    def make_node(self, x):
+        x = as_tensor_variable(x)
+        assert x.ndim == 2
+        sign = scalar(dtype=x.dtype)
+        det = scalar(dtype=x.dtype)
+        return Apply(self, [x], [sign, det])
+
+    def perform(self, node, inputs, outputs):
+        (x,) = inputs
+        (sign, det) = outputs
+        try:
+            sign[0], det[0] = (z.astype(x.dtype) for z in np.linalg.slogdet(x))
+        except Exception:
+            print("Failed to compute determinant", x)
+            raise
+
+    def infer_shape(self, fgraph, node, shapes):
+        return [(), ()]
+
+    def __str__(self):
+        return "SLogDet"
+
+
+slogdet = SLogDet()
+
+
 class Eig(Op):
     """
     Compute the eigenvalues and right eigenvectors of a square array.

tests/link/jax/test_nlinalg.py

@@ -85,6 +85,10 @@ def assert_fn(x, y):
     out_fg = FunctionGraph([x], outs)
     compare_jax_and_py(out_fg, [X.astype(config.floatX)], assert_fn=assert_fn)
 
+    outs = at_nlinalg.slogdet(x)
+    out_fg = FunctionGraph([x], outs)
+    compare_jax_and_py(out_fg, [X.astype(config.floatX)], assert_fn=assert_fn)
+
 
 @pytest.mark.xfail(
     version_parse(jax.__version__) >= version_parse("0.2.12"),

tests/link/numba/test_nlinalg.py

@@ -179,6 +179,41 @@ def test_Det(x, exc):
         )
 
 
+@pytest.mark.parametrize(
+    "x, exc",
+    [
+        (
+            set_test_value(
+                at.dmatrix(),
+                (lambda x: x.T.dot(x))(rng.random(size=(3, 3)).astype("float64")),
+            ),
+            None,
+        ),
+        (
+            set_test_value(
+                at.lmatrix(),
+                (lambda x: x.T.dot(x))(rng.poisson(size=(3, 3)).astype("int64")),
+            ),
+            None,
+        ),
+    ],
+)
+def test_SLogDet(x, exc):
+    g = nlinalg.SLogDet()(x)
+    g_fg = FunctionGraph(outputs=g)
+
+    cm = contextlib.suppress() if exc is None else pytest.warns(exc)
+    with cm:
+        compare_numba_and_py(
+            g_fg,
+            [
+                i.tag.test_value
+                for i in g_fg.inputs
+                if not isinstance(i, (SharedVariable, Constant))
+            ],
+        )
+
+
 # We were seeing some weird results in CI where the following two almost
 # sign-swapped results were being return from Numba and Python, respectively.
 # The issue might be related to https://github.com/numba/numba/issues/4519.

tests/tensor/test_extra_ops.py

@@ -1253,41 +1253,52 @@ def test_avoid_useless_subtensors(self):
     @pytest.mark.parametrize("linker", ["cvm", "py"])
     def test_perform(self, linker):
 
-        a = pytensor.shared(5)
+        a = pytensor.shared(np.full((3, 1, 1), 5))
+        s_0 = iscalar("s_0")
         s_1 = iscalar("s_1")
-        shape = (s_1, 1)
+        shape = (s_0, s_1, 1)
 
         bcast_res = broadcast_to(a, shape)
-        assert bcast_res.broadcastable == (False, True)
+        assert bcast_res.broadcastable == (False, False, True)
 
         bcast_fn = pytensor.function(
-            [s_1], bcast_res, mode=Mode(optimizer=None, linker=linker)
+            [s_0, s_1], bcast_res, mode=Mode(optimizer=None, linker=linker)
         )
         bcast_fn.vm.allow_gc = False
 
-        bcast_at = bcast_fn(4)
-        bcast_np = np.broadcast_to(5, (4, 1))
+        bcast_at = bcast_fn(3, 4)
+        bcast_np = np.broadcast_to(5, (3, 4, 1))
 
         assert np.array_equal(bcast_at, bcast_np)
 
-        bcast_var = bcast_fn.maker.fgraph.outputs[0].owner.inputs[0]
-        bcast_in = bcast_fn.vm.storage_map[a]
-        bcast_out = bcast_fn.vm.storage_map[bcast_var]
+        with pytest.raises(ValueError):
+            bcast_fn(5, 4)
 
         if linker != "py":
+            bcast_var = bcast_fn.maker.fgraph.outputs[0].owner.inputs[0]
+            bcast_in = bcast_fn.vm.storage_map[a]
+            bcast_out = bcast_fn.vm.storage_map[bcast_var]
             assert np.shares_memory(bcast_out[0], bcast_in[0])
 
+    def test_make_node_error_handling(self):
+        with pytest.raises(
+            ValueError,
+            match="Broadcast target shape has 1 dims, which is shorter than input with 2 dims",
+        ):
+            broadcast_to(at.zeros((3, 4)), (5,))
+
     @pytest.mark.skipif(
         not config.cxx, reason="G++ not available, so we need to skip this test."
     )
-    def test_memory_leak(self):
+    @pytest.mark.parametrize("valid", (True, False))
+    def test_memory_leak(self, valid):
         import gc
         import tracemalloc
 
         from pytensor.link.c.cvm import CVM
 
         n = 100_000
-        x = pytensor.shared(np.ones(n, dtype=np.float64))
+        x = pytensor.shared(np.ones((1, n), dtype=np.float64))
         y = broadcast_to(x, (5, n))
 
         f = pytensor.function([], y, mode=Mode(optimizer=None, linker="cvm"))
@@ -1303,8 +1314,17 @@ def test_memory_leak(self):
         blocks_last = None
         block_diffs = []
         for i in range(1, 50):
-            x.set_value(np.ones(n))
-            _ = f()
+            if valid:
+                x.set_value(np.ones((1, n)))
+                _ = f()
+            else:
+                x.set_value(np.ones((2, n)))
+                try:
+                    _ = f()
+                except ValueError:
+                    pass
+                else:
+                    raise RuntimeError("Should have failed")
             _ = gc.collect()
             blocks_i, _ = tracemalloc.get_traced_memory()
             if blocks_last is not None:
@@ -1313,7 +1333,7 @@ def test_memory_leak(self):
             blocks_last = blocks_i
 
         tracemalloc.stop()
-        assert np.allclose(np.mean(block_diffs), 0)
+        assert np.all(np.array(block_diffs) <= (0 + 1e-8))
 
     @pytest.mark.parametrize(
         "fn,input_dims",

tests/tensor/test_nlinalg.py

@@ -24,6 +24,7 @@
     norm,
     pinv,
     qr,
+    slogdet,
     svd,
     tensorinv,
     tensorsolve,
@@ -280,6 +281,18 @@ def test_det_shape():
     assert tuple(det_shape.data) == ()
 
 
+def test_slogdet():
+    rng = np.random.default_rng(utt.fetch_seed())
+
+    r = rng.standard_normal((5, 5)).astype(config.floatX)
+    x = matrix()
+    f = pytensor.function([x], slogdet(x))
+    f_sign, f_det = f(r)
+    sign, det = np.linalg.slogdet(r)
+    assert np.equal(sign, f_sign)
+    assert np.allclose(det, f_det)
+
+
 def test_trace():
     rng = np.random.default_rng(utt.fetch_seed())
     x = matrix()