Can a nnx.Module be in the carry of a jax.lax.fori_loop?

[wittenator]

I wanted to write a jax.lax.fori_loop that iteratively applies a model to some input with the model being in the carry argument. Weirdly enough, I get such an error:

TypeError                                 Traceback (most recent call last)

[<ipython-input-6-0c92d43c723b>](https://localhost:8080/#) in <cell line: 0>()
      4   return model
      5 
----> 6 jax.lax.fori_loop(0, 10, body, Counter())

    [... skipping hidden 5 frame]

[/usr/local/lib/python3.11/dist-packages/jax/_src/core.py](https://localhost:8080/#) in get_aval(x)
   1548   if hasattr(x, '__jax_array__'):
   1549     return get_aval(x.__jax_array__())
-> 1550   raise TypeError(f"Argument '{x}' of type '{typ}' is not a valid JAX type")
   1551 
   1552 get_ty = get_aval

TypeError: Argument 'Counter( # Count: 1 (4 B)
  count=Count( # 1 (4 B)
    value=Array(0, dtype=int32, weak_type=True)
  )
)' of type '<class '__main__.Counter'>' is not a valid JAX type

for this minimal reproducible example:

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

class Count(nnx.Variable): pass

class Counter(nnx.Module):
  def __init__(self):
    self.count = Count(jnp.array(0))

  def __call__(self):
    self.count += 1

@nnx.jit
def body(i, model):
  model()
  return model

jax.lax.fori_loop(0, 10, body, Counter())

I think this error is connected to the fact that Jax automatically compiles the body of a fori_loop via jax.jit, but even explicitly adding nnx.jit to the body_func does not resolve this error. Is there a better way to carry a model through a fori_loop?

There was another issue with nnx.Module and fori_loop some time ago, but the errors seem to be different: #4436

[wittenator]

Ah it seems that it is necessary to use nnx.fori_loop. My bad, I missed that transform