code: def test_remat_with_scan(self):
      class LinearBlock(nnx.Module):
          def __init__(self, *, rngs: nnx.Rngs):
              self.linear = nnx.Linear(3, 3, rngs=rngs)
  
          def __call__(
              self, x: jax.Array, _
          ) -> tp.Tuple[jax.Array, None]:
              x = self.linear(x)
              return x, None
  
      RematLinear = nnx.Remat(LinearBlock)
  
      ScanRematLinear = nnx.Scan(
          RematLinear,
          variable_axes={nnx.Param: 0},
          length=5,
      )
  
      m = ScanRematLinear(rngs=nnx.Rngs(0))
  
      assert m.scan_module.remat_module.linear.kernel.shape == (5, 3, 3)
      assert m.scan_module.remat_module.linear.bias.shape == (5, 3)
  
      y, _ = m(jnp.ones((1, 3)), None)
      assert y.shape == (1, 3)
  
      y, _ = m(jnp.ones((1, 3)), None)
      assert y.shape == (1, 3)
  
imports: import typing
  from functools import partial
  import jax
  import jax.numpy
  import pytest
  from flax.experimental import nnx
label_desc: Use image node_id 3 for calling the TestRemat obj's underlying member method code with example usage: obj.test_remat_with_scan() and returns: x,  None