code: def test_basic_remat(self):
      RematLinear = nnx.Remat(nnx.Linear)
  
      module = RematLinear(2, 3, rngs=nnx.Rngs(0))
  
      y = module(jnp.ones((1, 2)))
  
      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 1 for calling the TestRemat obj's underlying member method code with example usage: obj.test_basic_remat() without return types