code: def test_nnx_linen_equivalence(self, **kwargs):
      key = jax.random.key(42)
      rngs = nnx.Rngs(42)
      IN_FEATURES = 32
      OUT_FEATURES = 64
  
      x = jax.numpy.ones((1, IN_FEATURES))
      model_nnx = nnx.Linear.create_abstract(
          IN_FEATURES, OUT_FEATURES, **kwargs, rngs=rngs
      )
      model = linen.Dense(OUT_FEATURES, **kwargs)
      variables = model.init(key, x)
      model_nnx.kernel = variables["params"]["kernel"]
      if kwargs["use_bias"]:
          model_nnx.bias = variables["params"]["bias"]
  
      out_nnx = model_nnx(x)
      out = model.apply(variables, x)
      assert_array_equal(out, out_nnx)
  
file_path: flax/flax/experimental/nnx/tests/nn/test_linear.py
imports: import jax
  import jax.numpy
  from absl.testing import parameterized
  from jax.lax import Precision
  from numpy.testing import assert_array_equal
  from flax import linen
  from flax.experimental import nnx
label_desc: Use image node_id 1 for calling the TestLinenConsistency obj's underlying member method code with example usage: obj.test_nnx_linen_equivalence() without return types