code: def test_lookup(self, lookup_style, inp_ids):
      emb = embed.Embed(
          embedding_matrix=_EMBEDDING_MATRIX, lookup_style=lookup_style
      )
      np.testing.assert_allclose(
          emb(inp_ids),
          jnp.asarray(_EMBEDDING_MATRIX)[jnp.asarray(inp_ids)],
      )
      self.assertEqual(
          list(emb(inp_ids).shape),
          list(
              jnp.asarray(_EMBEDDING_MATRIX)[jnp.asarray(inp_ids)].shape
          ),
      )
  
file_path: dm-haiku/haiku/_src/embed_test.py
imports: import itertools
  from absl.testing import absltest
  from absl.testing import parameterized
  from haiku._src import embed
  from haiku._src import test_utils
  import jax.numpy
  import numpy
label_desc: Use image node_id 1 for calling the EmbedTest obj's underlying member method code with example usage: obj.test_lookup(lookup_style, inp_ids) without return types