nnx.remat (and other transformations) is easily misused.

nnx.remat needs to be applied to (unbound) functions, but it's easy to apply it to a bound method, e.g.

```
class Model(nnx.Module):

def __init__(self, *, rngs: nnx.Rngs):
self.linear = nnx.Linear(16, 32, rngs=rngs)
self.x_max = jnp.array(0.0)

def __call__(self, x):
return nnx.remat(self.block)(x)

def block(self, x):
self.x_max = jnp.maximum(self.x_max, x.max())
return self.linear(x)

model = Model(rngs=nnx.Rngs(0))
model(jnp.ones((3, 16)))
# >>> TraceContextError: Cannot mutate 'Model' from different trace level
```

Instead, we need to change `nnx.remat(self.block)(x)` to `nnx.remat(self.block.__func__)(self, x)`.

Is it possible to automate this? e.g., when the function passed to `nnx.remat` is bound, we automatically unbound it and pass the module as the first argument.