[BUG] qml.math.requires_grad returns True in cases where it shouldn't with jax.jacobian

[isaacdevlugt]

Expected behavior

Positional arguments specified in argnums should have qml.math.requires_grad == True, while those not listed in argnums should have qml.math.requires_grad == False.

Actual behavior

Positional arguments not listed in argnums have qml.math.requires_grad == True.

Additional information

No response

Source code

def f(x, y):
    print(qml.math.requires_grad(x), type(x))
    print(qml.math.requires_grad(y), type(y))
    return jnp.sin(x) * y ** 2
  
@partial(jax.jit, static_argnums=[2])
def g(x, y, argnums):
    return jax.jacobian(f, argnums=argnums)(x, y)

x = jnp.array([0.1, 0.2])
y = jnp.array([0.4, 0.3])

Tracebacks

>>> g(x, y, argnums=0)
True <class 'jax._src.interpreters.ad.JVPTracer'>
True <class 'jax._src.interpreters.partial_eval.DynamicJaxprTracer'>
Array([[0.15920068, 0.        ],
       [0.        , 0.088206  ]], dtype=float32)
>>> g(x, y, argnums=(0, 1))
True <class 'jax._src.interpreters.ad.JVPTracer'>
True <class 'jax._src.interpreters.ad.JVPTracer'>
(Array([[0.15920068, 0.        ],
        [0.        , 0.088206  ]], dtype=float32),
 Array([[0.07986674, 0.        ],
        [0.        , 0.1192016 ]], dtype=float32))

System information

Name: PennyLane
Version: 0.38.0
Summary: PennyLane is a cross-platform Python library for quantum computing, quantum machine learning, and quantum chemistry. Train a quantum computer the same way as a neural network.
Home-page: https://github.com/PennyLaneAI/pennylane
Author:
Author-email:
License: Apache License 2.0
Location: /Users/isaac/.virtualenvs/pennylane-stable/lib/python3.11/site-packages
Requires: appdirs, autograd, autoray, cachetools, networkx, numpy, packaging, pennylane-lightning, requests, rustworkx, scipy, toml, typing-extensions
Required-by: PennyLane-Catalyst, PennyLane_Lightning

Platform info:           macOS-14.6.1-arm64-arm-64bit
Python version:          3.11.9
Numpy version:           1.26.4
Scipy version:           1.12.0
Installed devices:
- nvidia.custatevec (PennyLane-Catalyst-0.8.0)
- nvidia.cutensornet (PennyLane-Catalyst-0.8.0)
- oqc.cloud (PennyLane-Catalyst-0.8.0)
- softwareq.qpp (PennyLane-Catalyst-0.8.0)
- lightning.qubit (PennyLane_Lightning-0.38.0)
- default.clifford (PennyLane-0.38.0)
- default.gaussian (PennyLane-0.38.0)
- default.mixed (PennyLane-0.38.0)
- default.qubit (PennyLane-0.38.0)
- default.qubit.autograd (PennyLane-0.38.0)
- default.qubit.jax (PennyLane-0.38.0)
- default.qubit.legacy (PennyLane-0.38.0)
- default.qubit.tf (PennyLane-0.38.0)
- default.qubit.torch (PennyLane-0.38.0)
- default.qutrit (PennyLane-0.38.0)
- default.qutrit.mixed (PennyLane-0.38.0)
- default.tensor (PennyLane-0.38.0)
- null.qubit (PennyLane-0.38.0)

Existing GitHub issues

• I have searched existing GitHub issues to make sure the issue does not already exist.