argnums parameter of catalyst.grad is ignored

[lasradorohan]

Parameter argnums of catalyst.grad() and catalyst.value_and_grad() is getting ignored. The gradients returned are of the first parameter regardless of argnums.

Code to reproduce:

import catalyst
import jax.numpy as jnp
import jax.random
import pennylane as qml
from pennylane import qjit, qnode

device = qml.device("lightning.qubit", wires=4)

@qjit
@qnode(device, interface="jax")
def circuit(inputs, weights):
    qml.AngleEmbedding(features=inputs, wires=range(4), rotation="X")
    for i in range(1, 4):
        qml.CRX(weights[i - 1], wires=[i, 0])
    return qml.expval(qml.PauliZ(wires=0))


if __name__ == '__main__':
    weights = jax.random.uniform(jax.random.key(0), shape=(3,), minval=0, maxval=jnp.pi)
    inputs = jax.random.uniform(jax.random.key(1), shape=(3,))
    print(catalyst.grad(circuit, argnums=0)(weights, inputs))
    print(catalyst.grad(circuit, argnums=1)(weights, inputs))
    print(catalyst.grad(circuit, argnums=(0, 1))(weights, inputs))
    print(catalyst.value_and_grad(circuit, argnums=0)(weights, inputs))
    print(catalyst.value_and_grad(circuit, argnums=1)(weights, inputs))
    print(catalyst.value_and_grad(circuit, argnums=(0, 1))(weights, inputs))

Output:

[-0.91515096 -0.0873711  -0.05631352]
[-0.91515096 -0.0873711  -0.05631352]
[-0.91515096 -0.0873711  -0.05631352]
(Array(-0.37879515, dtype=float64), Array([-0.91515096, -0.0873711 , -0.05631352], dtype=float64))
(Array(-0.37879515, dtype=float64), Array([-0.91515096, -0.0873711 , -0.05631352], dtype=float64))
(Array(-0.37879515, dtype=float64), Array([-0.91515096, -0.0873711 , -0.05631352], dtype=float64))

Issue with output: all the gradients are from the first input and therefore the same.

[lasradorohan]

PR: #1478