Avoid duplicated inputs in KroneckerProduct OpFromGraph

Author: ricardoV94

pytensor/compile/builders.py

@@ -400,6 +400,15 @@ def __init__(
             Check :func:`pytensor.function` for more arguments, only works when not
             inline.
         """
+        ignore_unused_inputs = kwargs.get("on_unused_input", False) == "ignore"
+        if not ignore_unused_inputs and len(inputs) != len(set(inputs)):
+            var_counts = {var: inputs.count(var) for var in inputs}
+            duplicated_inputs = [var for var, count in var_counts.items() if count > 1]
+            raise ValueError(
+                f"There following variables were provided more than once as inputs to the OpFromGraph, resulting in an "
+                f"invalid graph: {duplicated_inputs}. Use dummy variables or var.copy() to distinguish "
+                f"variables when creating the OpFromGraph graph."
+            )
 
         if not (isinstance(inputs, list) and isinstance(outputs, list)):
             raise TypeError("Inputs and outputs must be lists")

pytensor/tensor/nlinalg.py

@@ -1034,6 +1034,11 @@ def kron(a, b):
     """
     a = as_tensor_variable(a)
     b = as_tensor_variable(b)
+
+    if a is b:
+        # In case a is the same as b, we need a different variable to build the OFG
+        b = a.copy()
+
     if a.ndim + b.ndim <= 2:
         raise TypeError(
             "kron: inputs dimensions must sum to 3 or more. "

tests/compile/test_builders.py

@@ -118,6 +118,7 @@ def test_grad_grad(self, cls_ofg):
         f = op(x, y, z)
         f = f - grad(pt_sum(f), y)
         f = f - grad(pt_sum(f), y)
+
         fn = function([x, y, z], f)
         xv = np.ones((2, 2), dtype=config.floatX)
         yv = np.ones((2, 2), dtype=config.floatX) * 3
@@ -584,6 +585,22 @@ def test_explicit_input_from_shared(self):
         out = test_ofg(y, y)
         assert out.eval() == 4
 
+    def test_repeated_inputs(self):
+        x = pt.dscalar("x")
+        y = pt.dscalar("y")
+
+        with pytest.raises(
+            ValueError,
+            match="There following variables were provided more than once as inputs to the "
+            "OpFromGraph",
+        ):
+            OpFromGraph([x, x], [x + y])
+
+        # Test that repeated inputs will be allowed if unused inputs are ignored
+        g = OpFromGraph([x, x, y], [x + y], on_unused_input="ignore")
+        f = g(x, x, y)
+        assert f.eval({x: 5, y: 5}) == 10
+
 
 @config.change_flags(floatX="float64")
 def test_debugprint():

tests/tensor/test_slinalg.py

@@ -514,8 +514,8 @@ def test_expm_grad_3():
 def test_solve_discrete_lyapunov_via_direct_real():
     N = 5
     rng = np.random.default_rng(utt.fetch_seed())
-    a = pt.dmatrix()
-    q = pt.dmatrix()
+    a = pt.dmatrix("a")
+    q = pt.dmatrix("q")
     f = function([a, q], [solve_discrete_lyapunov(a, q, method="direct")])
 
     A = rng.normal(size=(N, N))