Taking K-gradient steps with NNX

[pecey]

Hi. I was wondering if there is a way to take k-gradient steps with a jax.fori_loop. I can do it with a for loop, but as k can be large, I wanted to embed it with a fori_loop. Below is an example of what I have tried but was not able to get it to work. I get a Cannot mutate OptArray from a different trace level with a link to TraceContextError, but I couldn't find anything helpful.

Also, if I have two variables in my model and I want to optimize them separately, how do I go about that? For example: here I have self.W and a separate optimizer for it called w_optim, and I tried using filters to ensure they acted on the correct variables, but as soon as I uncomment them, I get Mismatch custom node data: ('A', 'W') != ('A',) while updating the gradients.

from flax import nnx
import jax.numpy as jnp
import jax
import optax


class Agent(nnx.Module):
    def __init__(self, K, rngs: nnx.Rngs):
        self.A = nnx.Param(jax.random.normal(rngs.params(), (K,)))
        # self.W = nnx.Param(jax.random.normal(rngs.params(), (K,)))

agent = Agent(4, nnx.Rngs(0))
a_optim = nnx.Optimizer(agent, optax.adam(1e-2), wrt=nnx.All(nnx.Param, nnx.PathContains("A")))
# w_optim = nnx.Optimizer(agent, optax.sgd(1e-2), wrt=nnx.All(nnx.Param, nnx.PathContains("W")))


def compute_loss(A, B, W):
     return jnp.sum(W * (A - B)**2) 

def update_a(agent, B, W, optimizer, n=500):
    def loss_fn(model:Agent):
        return compute_loss(model.A, B, W)

    def update_step(i, val):
        agent, loss = val
        loss, grads = nnx.value_and_grad(loss_fn)(agent)
        optimizer.update(agent, grads)
        return agent, loss
    agent, loss = jax.lax.fori_loop(0, n, update_step, (agent, 0.0))
    
    # for _ in range(n):
    #     loss, grads = nnx.value_and_grad(loss_fn)(agent)
    #     optimizer.update(agent, grads)
    return agent, loss


A = jnp.array([1. ,  0.85 ,0.8 , 0.7])
B = jnp.array([0.35268408, 1.2938553,  1.4211403,  0.892035])
W = jnp.array([0.3061288 , 0.27737275, 0.0714835 ,0.34501496])

update_a(agent, B, W, a_optim, n = 10)

I am not sure what is it that I am doing wrong. Any help will be greatly appreciated. Thank you.