Accessing intermediary computations of remat_scan

[organic-chemistry]

Hello, I have a NN that use remat_scan to perform a loop of length 100, and after training I would like to be able to
access the intermediary steps.
Here is a toy example of what I would like to do.

from flax import linen as nn
import jax.numpy as jnp
import jax

class BigModel(nn.Module):
    
    def setup(self):
        features = 8
        self.DenseStack = nn.remat_scan(nn.Dense, lengths=(10, 10))(features)
        
    @nn.compact
    def __call__(self, x):
        return self.DenseStack(x)
    
    def history(self, x ,train=False):
        self.DenseStack.lengths=(1,1)
        ret = []
        for t in range(100):
            x = self.DenseStack(x)
            ret.append(x)
        return ret
    
model = BigModel()
X_train =  jnp.zeros([32,8])+1
variables = model.init(jax.random.PRNGKey(0),X_train)
output = model.apply(variables, X_train) 
output = model.apply(variables, X_train,method=model.history) 

so this throw this a SetAttributeFrozenModuleError .
I understand, but I have not been able to reuse or acces the nn.Dense module that use the DenseStack scan.

Best!