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..9fe791d8ca 100644 --- a/tests/link/jax/test_wrap_jax.py +++ b/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,))