Fixed incorrect use of walrus operator

[lasradorohan]

Incorrect use of the walrus operator caused argnums to be assigned a boolean instead of the value of rhs value.

Context:
Parameter argnums of catalyst.grad() and catalyst.value_and_grad() was getting ignored. The gradients returned were of the first parameter regardless of argnums.
In incorrect use of the walrus operator := was found in the code.

Description of the Change:
Brackets were added to ensure correct assignment of value.

Benefits:
argnums now works as expected.

Possible Drawbacks:

Related GitHub Issues:
#1477

[Qottmann]

I dont think this is actually a bug. Also note that pylint complains when you do these brackets around the walrus operator.

I personally also prefer it with the brackets for better readability though, so wouldnt mind the change. But note that there are several occurences of this pattern in the codebase, so it wouldnt be just changing this line here.

nevermind this is a different case here than I thought

[dime10]

Hi @lasradorohan, thank you so much for uncovering this issue and even providing a fix for it 🚀 🚀

[dime10]

I've approved the CI to run on your PR to see if anything breaks (it will also check formatting and linting). In order to merge this PR, we'll require the following:

• Normally all bug fixes require one (or more) new tests that would have failed before and now pass. This is to prevent us from regressing into the bug again in the future. If you would like to add a test case (maybe a slimmed down version of how you originally discovered it) you can add a test in this file. Otherwise, let me know and we can add a test to your PR.
• Update the changelog saying what was fixed with this PR in this file, and most importantly, don't forget to add your name to the contributor list at the bottom of the file :)

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 96.70%. Comparing base (78c7251) to head (015ea61).
Report is 4 commits behind head on main.

Additional details and impacted files

@@           Coverage Diff           @@
##             main    #1478   +/-   ##
=======================================
  Coverage   96.70%   96.70%           
=======================================
  Files          76       76           
  Lines        8173     8173           
  Branches      846      846           
=======================================
  Hits         7904     7904           
  Misses        215      215           
  Partials       54       54           

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

[dime10]

Nice work, happy to approve 💯 (looks like you just need to run black frontend to satisfy the CI)

[dime10]

I don't think we need to test without qjit, since we would be testing PennyLane in that case (which has its own test suite):

Suggested change

	@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)
	def test_grad_argnums(argnums):
	"""Tests https://github.com/PennyLaneAI/catalyst/issues/1477"""
	@qjit
	@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))

[dime10]

Ideally we might want to compare against a computation not involving Catalyst at all, which can be achieved like this (applying jax.grad function is a sort of Catalyst/JAX hybrid):

Suggested change

	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)
	expected = jax.grad(circuit.original_function, 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.original_function, argnums=argnums)(weights, inputs)