Treat discrete JAXOp inputs as disconnected in the gradient

pytensor/link/jax/ops.py

@@ -10,7 +10,7 @@
 from pytensor.gradient import DisconnectedType
 from pytensor.graph import Apply, Op, Variable
 from pytensor.tensor.basic import as_tensor, infer_static_shape
-from pytensor.tensor.type import TensorType
+from pytensor.tensor.type import TensorType, discrete_dtypes
 
 
 class JAXOp(Op):
@@ -135,6 +135,13 @@ def perform_jax(self, *inputs):
             return outputs[0]
         return outputs
 
+    def connection_pattern(self, node):
+        """Mark discrete inputs as disconnected from every output."""
+        return [
+            [input_type.dtype not in discrete_dtypes] * len(self.output_types)
+            for input_type in self.input_types
+        ]
+
     def pullback(self, inputs, outputs, output_gradients):
         """Compute gradients using JAX's vector-Jacobian product (VJP)."""
         import jax
@@ -146,6 +153,15 @@ def pullback(self, inputs, outputs, output_gradients):
             if not isinstance(output_grad.type, DisconnectedType)
         ]
 
+        # Integer and boolean inputs are not differentiable. JAX gives them
+        # float0 cotangents, which have no PyTensor equivalent, so they are
+        # held constant in the VJP and reported as disconnected.
+        differentiable_input_indices = [
+            i
+            for i, input_type in enumerate(self.input_types)
+            if input_type.dtype not in discrete_dtypes
+        ]
+
         num_inputs = len(inputs)
 
         def vjp_operation(*args):
@@ -154,15 +170,23 @@ def vjp_operation(*args):
             cotangent_vectors = args[num_inputs:]
             assert len(cotangent_vectors) == len(connected_output_indices)
 
-            def restricted_function(*input_values):
-                """Restricted function that only returns connected outputs."""
-                outputs = self.jax_func(*input_values)
+            def restricted_function(*differentiable_values):
+                """Restricted function of the differentiable inputs, returning connected outputs."""
+                all_input_values = list(input_values)
+                for i, value in zip(
+                    differentiable_input_indices, differentiable_values, strict=True
+                ):
+                    all_input_values[i] = value
+                outputs = self.jax_func(*all_input_values)
                 return [
                     outputs[i].astype(self.output_types[i].dtype)
                     for i in connected_output_indices
                 ]
 
-            _primals, vjp_function = jax.vjp(restricted_function, *input_values)
+            _primals, vjp_function = jax.vjp(
+                restricted_function,
+                *(input_values[i] for i in differentiable_input_indices),
+            )
             output_dtypes = [
                 self.output_types[i].dtype for i in connected_output_indices
             ]
@@ -184,15 +208,21 @@ def restricted_function(*input_values):
         vjp_op = JAXOp(
             self.input_types
             + tuple(self.output_types[i] for i in connected_output_indices),
-            [self.input_types[i] for i in range(num_inputs)],
+            [self.input_types[i] for i in differentiable_input_indices],
             vjp_operation,
             name=name,
         )
 
-        return vjp_op(
+        differentiable_input_gradients = vjp_op(
             *[*inputs, *[output_gradients[i] for i in connected_output_indices]],
             return_list=True,
         )
+        input_gradients = [DisconnectedType()() for _ in inputs]
+        for i, input_gradient in zip(
+            differentiable_input_indices, differentiable_input_gradients, strict=True
+        ):
+            input_gradients[i] = input_gradient
+        return input_gradients
 
 
 def wrap_jax(jax_function=None, *, allow_eval=True):

tests/link/jax/test_wrap_jax.py

@@ -3,6 +3,7 @@
 
 from pytensor import config, grad, wrap_jax
 from pytensor.compile.sharedvalue import shared
+from pytensor.gradient import DisconnectedInputError
 from pytensor.link.jax.ops import JAXOp
 from pytensor.scalar import all_types
 from pytensor.tensor import TensorType, tensor
@@ -351,6 +352,69 @@ def f(x, y):
     compare_jax_and_py([x, y], [out[1], *grad_out], test_values)
 
 
+def test_discrete_input():
+    # Integer inputs are not differentiable, so they must be reported as
+    # disconnected instead of getting an integer gradient. Regression test for #2072.
+    rng = np.random.default_rng(12)
+    x = tensor("x", shape=(5,))
+    idx = tensor("idx", shape=(3,), dtype="int32")
+    test_values = [
+        rng.normal(size=x.type.shape).astype(config.floatX),
+        np.array([0, 2, 2], dtype="int32"),
+    ]
+
+    def f(x, idx):
+        return jax.numpy.sum(x[idx])
+
+    out = wrap_jax(f)(x, idx)
+    [grad_out] = grad(out, [x])
+    _, jax_res = compare_jax_and_py([x, idx], [out, grad_out], test_values)
+    np.testing.assert_allclose(jax_res[1], [1.0, 0.0, 2.0, 0.0, 0.0])
+
+    with pytest.raises(DisconnectedInputError):
+        grad(out, [idx])
+
+
+def test_mixed_input_types():
+    # Discrete inputs must stay disconnected when they are interleaved with
+    # differentiable ones, and the gradients of the latter must be unaffected.
+    rng = np.random.default_rng(13)
+    x = tensor("x", shape=(5,))
+    idx = tensor("idx", shape=(3,), dtype="int32")
+    y = tensor("y", shape=(3,))
+    mask = tensor("mask", shape=(3,), dtype="bool")
+    x_test = rng.normal(size=(5,)).astype(config.floatX)
+    idx_test = np.array([0, 2, 2], dtype="int32")
+    y_test = rng.normal(size=(3,)).astype(config.floatX)
+    mask_test = np.array([True, False, True])
+    test_values = [x_test, idx_test, y_test, mask_test]
+
+    def f(x, idx, y, mask):
+        return jax.numpy.sum(x[idx] * y * mask)
+
+    out = wrap_jax(f)(x, idx, y, mask)
+    assert out.owner.op.connection_pattern(out.owner) == [
+        [True],
+        [False],
+        [True],
+        [False],
+    ]
+
+    grad_x, grad_y = grad(out, [x, y])
+    _, jax_res = compare_jax_and_py(
+        [x, idx, y, mask], [out, grad_x, grad_y], test_values
+    )
+
+    expected_grad_x = np.zeros_like(x_test)
+    np.add.at(expected_grad_x, idx_test, y_test * mask_test)
+    np.testing.assert_allclose(jax_res[1], expected_grad_x)
+    np.testing.assert_allclose(jax_res[2], x_test[idx_test] * mask_test)
+
+    for discrete_input in (idx, mask):
+        with pytest.raises(DisconnectedInputError):
+            grad(out, [discrete_input])
+
+
 def test_unknown_static_shape():
     rng = np.random.default_rng(11)
     x = tensor("x", shape=(3,))