code: def test_nnx_linen_equivalence(self, **kwargs):
      key = jax.random.key(42)
      rngs = nnx.Rngs(42)
      IN_FEATURES = 3
      OUT_FEATURES = 6
      INPUT_SHAPE = (24, 9, IN_FEATURES)
      kwargs["kernel_size"] = (7, 4)
  
      # Cannot use string padding specification for transpose conv
      if (
          isinstance(kwargs["input_dilation"], Sequence)
          or kwargs["input_dilation"] > 1
      ):
          kwargs["padding"] = (4, 2)
  
      x = jax.numpy.ones(INPUT_SHAPE)
      model_nnx = nnx.Conv.create_abstract(
          IN_FEATURES, OUT_FEATURES, **kwargs, rngs=rngs
      )
      model = linen.Conv(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_conv.py
imports: from collections.abc import Sequence
  import jax
  from absl.testing import parameterized
  from jax import numpy
  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 TestConvLinenConsistency obj's underlying member method code with example usage: obj.test_nnx_linen_equivalence() without return types