Treat discrete JAXOp inputs as disconnected in the gradient

[rishabhsai]

Description

JAXOp.pullback sent every input through jax.vjp and declared the VJP op's output types as [self.input_types[i] for i in range(num_inputs)], so an integer input came back with an integer gradient and pytensor/gradient.py rejected it with "grad illegally returned an integer-valued variable". Returning DisconnectedType instead was not possible either, because JAXOp had no connection_pattern, so Blockwise assumed every input was connected and refused the disconnected term. That makes grad fail for any wrap_jax function taking an index array, with or without pt.vectorize.

Discrete inputs are now held constant in the VJP, jax.vjp only sees the float inputs, and their gradient is reported as DisconnectedType. JAXOp.connection_pattern marks discrete inputs as disconnected from every output, which Blockwise.connection_pattern picks up for the vectorized case. This is the same arrangement AdvancedSubtensor uses for its index inputs.

Behaviour change: grad with respect to a discrete wrap_jax input now raises DisconnectedInputError rather than silently producing an integer zero.

Related Issue

• Closes BUG: JAXOp.grad fails when inputs include integer-typed tensors #2072
• Related to #

Checklist

• Checked that the pre-commit linting/style checks pass
• Included tests that prove the fix is effective or that the new feature works
• Added necessary documentation (docstrings and/or example notebooks)
• If you are a pro: each commit corresponds to a relevant logical change

Type of change

• New feature / enhancement
• Bug fix
• Documentation
• Maintenance
• Other (please specify):

[ricardoV94]

this is not the same, you now used a continuous input? I'd also like to see a test for wrap_jax with mixed input types. The blockwise itself was a symptom downstream, I don't think we need to test it

[rishabhsai]

Dropped the Blockwise block, agreed it was a downstream symptom. On the continuous input: the discrete assertions were the two lines above it (int32 idx, DisconnectedInputError), and the vectorized block did differentiate only the float batched_x, so it was carrying the connection_pattern coverage by accident rather than asserting anything discrete. Replaced it with test_mixed_input_types, which interleaves float x, int32 idx, float y and bool mask, checks the numeric gradients of x and y, asserts both discrete inputs raise DisconnectedInputError, and asserts connection_pattern on the op directly so that half of the fix stays covered without Blockwise.