Fixed incorrect use of walrus operator

doc/releases/changelog-dev.md

@@ -14,6 +14,9 @@
 
 <h3>Bug fixes 🐛</h3>
 
+* Fixed `argnums` parameter of `grad` and `value_and_grad` being ignored.
+  [(#1478)](https://github.com/PennyLaneAI/catalyst/pull/1478)
+
 <h3>Internal changes ⚙️</h3>
 
 * `from_plxpr` now uses the `qml.capture.PlxprInterpreter` class for reduced code duplication.
@@ -48,4 +51,5 @@ This release contains contributions from (in alphabetical order):
 Sengthai Heng,
 Christina Lee,
 Mehrdad Malekmohammadi,
-Paul Haochen Wang.
+Paul Haochen Wang,
+Rohan Nolan Lasrado.

frontend/catalyst/api_extensions/differentiation.py

@@ -696,7 +696,7 @@ def __call__(self, *args, **kwargs):
                         results, in_arg_tree, out_tree, grad_params, len(jaxpr.out_avals)
                     )
             else:
-                if argnums := self.grad_params.argnums is None:
+                if (argnums := self.grad_params.argnums) is None:
                     argnums = 0
                 if self.grad_params.scalar_out:
                     if self.grad_params.with_value:

frontend/test/pytest/test_gradient.py

@@ -19,7 +19,7 @@
 import pennylane as qml
 import pytest
 from jax import numpy as jnp
-from jax.tree_util import tree_flatten
+from jax.tree_util import tree_flatten, tree_map, tree_all, tree_structure
 
 import catalyst.utils.calculate_grad_shape as infer
 from catalyst import (
@@ -1780,5 +1780,35 @@ def fn(x):
         assert np.allclose(res_pattern_partial, expected)
 
 
+@pytest.mark.parametrize("argnums", [0, 1, (0, 1)])
+@pytest.mark.parametrize("transform_qjit", [False, True])
+def test_grad_argnums(argnums, transform_qjit):
+    """Tests https://github.com/PennyLaneAI/catalyst/issues/1477"""
+    @qml.qnode(device=qml.device("lightning.qubit", wires=4), 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 transform_qjit:
+        circuit = qjit(circuit)
+
+    weights = jnp.array([3.0326467, 0.98860157, 1.9887222])
+    inputs = jnp.array([0.9653214, 0.31468165, 0.63302994])
+
+    def compare_structure_and_value(o1, o2):
+        return tree_structure(o1) == tree_structure(o2) and \
+            tree_all(tree_map(jnp.allclose, o1, o2))
+
+    result = grad(circuit, argnums=argnums)(weights, inputs)
+    expected = jax.grad(circuit, argnums=argnums)(weights, inputs)
+    assert compare_structure_and_value(result, expected)
+
+    _, result = value_and_grad(circuit, argnums=argnums)(weights, inputs)
+    _, expected = jax.value_and_grad(circuit, argnums=argnums)(weights, inputs)
+    assert compare_structure_and_value(result, expected)
+
+
 if __name__ == "__main__":
     pytest.main(["-x", __file__])