[nnx] grad accepts argnums

[cgarciae]

What does this PR do?

nnx.grad now accepts argnums and multiple graph nodes can be passed.

[codecov-commenter]

Codecov Report

Attention: Patch coverage is 93.42105% with 5 lines in your changes are missing coverage. Please review.

Project coverage is 60.51%. Comparing base (cc740d4) to head (3bebf97).

Files	Patch %	Lines
flax/experimental/nnx/nnx/transforms.py	83.33%	5 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #3798      +/-   ##
==========================================
+ Coverage   60.34%   60.51%   +0.17%     
==========================================
  Files         101      101              
  Lines       12862    12908      +46     
==========================================
+ Hits         7761     7811      +50     
+ Misses       5101     5097       -4     

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

[chiamp]

Suggested change

	def test_multiple_graph_nodes(self):
	rngs = nnx.Rngs(0)
	m1 = nnx.Linear(2, 3, rngs=rngs)
	m2 = nnx.Linear(3, 3, rngs=rngs)
	loss_fn = lambda m1, m2, x, y: jnp.mean((m2(m1(x)) - y) ** 2)
	grad_fn = nnx.grad(loss_fn, argnums=(0, 1), wrt=nnx.Param)
	x = jax.random.uniform(rngs(), (1, 2))
	y = jnp.ones((1, 3))
	grads_m1, grads_m2 = grad_fn(m1, m2, x, y)
	assert 'kernel' in grads_m1
	assert grads_m1.kernel.raw_value.shape == (2, 3)
	assert 'bias' in grads_m1
	assert grads_m1.bias.raw_value.shape == (3,)
	assert 'kernel' in grads_m2
	assert grads_m2.kernel.raw_value.shape == (3, 3)
	assert 'bias' in grads_m2
	assert grads_m2.bias.raw_value.shape == (3,)
	def test_multiple_graph_nodes_mix_positions(self):
	rngs = nnx.Rngs(0)
	m1 = nnx.Linear(2, 3, rngs=rngs)
	m2 = nnx.Linear(3, 3, rngs=rngs)
	loss_fn = lambda x, m1, y, m2: jnp.mean((m2(m1(x)) - y) ** 2)
	grad_fn = nnx.grad(loss_fn, argnums=(1, 3), wrt=nnx.Param)
	x = jax.random.uniform(rngs(), (1, 2))
	y = jnp.ones((1, 3))
	grads_m1, grads_m2 = grad_fn(x, m1, y, m2)
	assert 'kernel' in grads_m1
	assert grads_m1.kernel.raw_value.shape == (2, 3)
	assert 'bias' in grads_m1
	assert grads_m1.bias.raw_value.shape == (3,)
	assert 'kernel' in grads_m2
	assert grads_m2.kernel.raw_value.shape == (3, 3)
	assert 'bias' in grads_m2
	assert grads_m2.bias.raw_value.shape == (3,)
	@parameterized.parameters(
	{'loss_fn': lambda m1, m2, x, y: jnp.mean((m2(m1(x)) - y) ** 2), 'argnums': (0, 1)},
	{'loss_fn': lambda x, m1, y, m2: jnp.mean((m2(m1(x)) - y) ** 2), 'argnums': (1, 3)},
	)
	def test_multiple_graph_nodes(self, loss_fn, argnums):
	rngs = nnx.Rngs(0)
	m1 = nnx.Linear(2, 3, rngs=rngs)
	m2 = nnx.Linear(3, 3, rngs=rngs)
	grad_fn = nnx.grad(loss_fn, argnums=argnums, wrt=nnx.Param)
	x = jax.random.uniform(rngs(), (1, 2))
	y = jnp.ones((1, 3))
	grads_m1, grads_m2 = grad_fn(m1, m2, x, y)
	assert 'kernel' in grads_m1
	assert grads_m1.kernel.raw_value.shape == (2, 3)
	assert 'bias' in grads_m1
	assert grads_m1.bias.raw_value.shape == (3,)
	assert 'kernel' in grads_m2
	assert grads_m2.kernel.raw_value.shape == (3, 3)
	assert 'bias' in grads_m2
	assert grads_m2.bias.raw_value.shape == (3,)

The two tests seem similar enough that we could combine them using parameterized.parameters

[cgarciae]

I think you can do it but the suggested implementation is not quite right, you also need to parametrize argnums and loss_fn is being redefined.

[chiamp]

ah yes good catch, I fixed it now

[cgarciae]

still doesn't work because you need to pass arguments to grad_fn in different positions.

[cgarciae]

merging for now