Using flax.linen.intercept_methods on an nnx module

Hi! I want to add a hook / interceptor to an nnx module (e.g. that prints the shapes of the input and output). This is straightforward to do for a module defined via the old flax.linen api (as shown [here](https://flax-linen.readthedocs.io/en/latest/api_reference/flax.linen/module.html#flax.linen.intercept_methods) in the docs):

```python
import flax.linen as nn
import jax.numpy as jnp

class Foo(nn.Module):
def __call__(self, x):
return x

def my_interceptor1(next_fun, args, kwargs, context):
print('calling my_interceptor1')
return next_fun(*args, **kwargs)

foo = Foo()
with nn.intercept_methods(my_interceptor1):
_ = foo(jnp.ones([1]))
# >> calling my_interceptor1
```

However, this does not seem to work for an nnx module:

```python
from flax import nnx # Using nnx instead of flax.linen
import jax.numpy as jnp

class Foo(nnx.Module): # changed to nnx.Module
def __call__(self, x):
return x

def my_interceptor1(next_fun, args, kwargs, context):
print('calling my_interceptor1')
return next_fun(*args, **kwargs)

foo = Foo()
with nn.intercept_methods(my_interceptor1):
_ = foo(jnp.ones([1]))
# No output
```

I'm assuming I should *not* assume the linen interceptor to work for an nnx module.

My questions are:

1. Is there a way to add an interceptor / hook to an nnx module? And if not:
2. Are there any plans to add one by either:

a. Migrating the old implementation?
b. Implementing a new one?

I would be happy to help out with either one if I could get some pointers :)