code: def test_nnx_linen_equivalence(
      self, input_dtype, num_embeddings, **kwargs
  ):
      key = jax.random.key(42)
      rngs = nnx.Rngs(42)
      IN_FEATURES = 32
      NUM_EMBEDDINGS = num_embeddings
  
      x = jax.numpy.arange(NUM_EMBEDDINGS, dtype=input_dtype)
      model_nnx = nnx.Embed.create_abstract(
          NUM_EMBEDDINGS, IN_FEATURES, **kwargs, rngs=rngs
      )
      model = linen.Embed(NUM_EMBEDDINGS, IN_FEATURES, **kwargs)
      variables = model.init(key, x)
      model_nnx.embedding = variables["params"]["embedding"]
  
      out_nnx = model_nnx(x)
      out = model.apply(variables, x)
      assert_array_equal(out, out_nnx)
  
      x = jax.numpy.ones((10,), dtype=input_dtype) * 10
      out_nnx = model_nnx(x)
      out = model.apply(variables, x)
      assert_array_equal(out, out_nnx)
      assert_array_equal(
          jax.numpy.isnan(out).all(), jax.numpy.array([True])
      )
  
file_path: flax/flax/experimental/nnx/tests/nn/test_embed.py
imports: import jax
  from absl.testing import parameterized
  from jax import numpy
  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(input_dtype, num_embeddings) without return types