From dee1ae2cab5be87cb022ced6248190c3a8c02504 Mon Sep 17 00:00:00 2001 From: Rishabh Date: Mon, 14 Sep 2026 19:30:46 -0400 Subject: [PATCH 1/2] Treat discrete JAXOp inputs as disconnected in the gradient --- pytensor/link/jax/ops.py | 44 +++++++++++++++++++++++++++------ tests/link/jax/test_wrap_jax.py | 37 ++++++++++++++++++++++++++- 2 files changed, 73 insertions(+), 8 deletions(-) diff --git a/pytensor/link/jax/ops.py b/pytensor/link/jax/ops.py index db2f7d8c31..7ef98df143 100644 --- a/pytensor/link/jax/ops.py +++ b/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): diff --git a/tests/link/jax/test_wrap_jax.py b/tests/link/jax/test_wrap_jax.py index b328ab22c5..b8b094fcb4 100644 --- a/tests/link/jax/test_wrap_jax.py +++ b/tests/link/jax/test_wrap_jax.py @@ -3,9 +3,10 @@ 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 +from pytensor.tensor import TensorType, tensor, vectorize from tests.link.jax.test_basic import compare_jax_and_py @@ -351,6 +352,40 @@ 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]) + + # Blockwise consults the core Op's connection_pattern, so the same must hold + # when the wrapped function is vectorized. + batched_x = tensor("batched_x", shape=(2, 5)) + batched_out = vectorize(wrap_jax(f), signature="(n),(m)->()")(batched_x, idx) + [batched_grad_out] = grad(batched_out.sum(), [batched_x]) + compare_jax_and_py( + [batched_x, idx], + [batched_out, batched_grad_out], + [np.repeat(test_values[0][None], 2, axis=0), test_values[1]], + ) + + def test_unknown_static_shape(): rng = np.random.default_rng(11) x = tensor("x", shape=(3,)) From 7c85996ed17a9a98701e57916e7eb2f845be0f13 Mon Sep 17 00:00:00 2001 From: Rishabh Date: Tue, 15 Sep 2026 23:55:14 -0400 Subject: [PATCH 2/2] Test wrap_jax gradients with mixed discrete and float inputs --- tests/link/jax/test_wrap_jax.py | 49 ++++++++++++++++++++++++++------- 1 file changed, 39 insertions(+), 10 deletions(-) diff --git a/tests/link/jax/test_wrap_jax.py b/tests/link/jax/test_wrap_jax.py index b8b094fcb4..9fe791d8ca 100644 --- a/tests/link/jax/test_wrap_jax.py +++ b/tests/link/jax/test_wrap_jax.py @@ -6,7 +6,7 @@ from pytensor.gradient import DisconnectedInputError from pytensor.link.jax.ops import JAXOp from pytensor.scalar import all_types -from pytensor.tensor import TensorType, tensor, vectorize +from pytensor.tensor import TensorType, tensor from tests.link.jax.test_basic import compare_jax_and_py @@ -374,17 +374,46 @@ def f(x, idx): with pytest.raises(DisconnectedInputError): grad(out, [idx]) - # Blockwise consults the core Op's connection_pattern, so the same must hold - # when the wrapped function is vectorized. - batched_x = tensor("batched_x", shape=(2, 5)) - batched_out = vectorize(wrap_jax(f), signature="(n),(m)->()")(batched_x, idx) - [batched_grad_out] = grad(batched_out.sum(), [batched_x]) - compare_jax_and_py( - [batched_x, idx], - [batched_out, batched_grad_out], - [np.repeat(test_values[0][None], 2, axis=0), test_values[1]], + +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)