class_docstr: Module for embedding tokens in a low-dimensional space.
class_docstr_tok: ["Module","for","embedding","tokens","in","a","low-dimensional","space","."]
docstr: Lookup embeddings.
  
  Looks up an embedding vector for each value in ``ids``. All ids must be
  within ``[0, vocab_size)`` to prevent ``NaN``\ s from propagating.
  
  Args:
    ids: integer array.
    lookup_style: Overrides the ``lookup_style`` given in the constructor.
    precision: Overrides the ``precision`` given in the constructor.
  
  Returns:
    Tensor of ``ids.shape + [embedding_dim]``.
  
  Raises:
    AttributeError: If ``lookup_style`` is not valid.
    ValueError: If ``ids`` is not an integer array.
docstr_tok: ["Lookup","embeddings",".","Looks","up","an","embedding","vector","for","each","value","in","``","ids","``",".","All","ids","must","be","within","``","[","0",",","vocab_size",")","``","to","prevent","``","NaN","``","\\","s","from","propagating",".","Args",":","ids",":","integer","array",".","lookup_style",":","Overrides","the","``","lookup_style","``","given","in","the","constructor",".","precision",":","Overrides","the","``","precision","``","given","in","the","constructor",".","Returns",":","Tensor","of","``","ids.shape","+","[","embedding_dim","]","``",".","Raises",":","AttributeError",":","If","``","lookup_style","``","is","not","valid",".","ValueError",":","If","``","ids","``","is","not","an","integer","array","."]
code: def __call__(
      self,
      ids: Union[jax.Array, Sequence[int]],
      lookup_style: Optional[Union[str, hk.EmbedLookupStyle]] = None,
      precision: Optional[jax.lax.Precision] = None,
  ) -> jax.Array:
      r"""Lookup embeddings.
  
      Looks up an embedding vector for each value in ``ids``. All ids must be
      within ``[0, vocab_size)`` to prevent ``NaN``\ s from propagating.
  
      Args:
        ids: integer array.
        lookup_style: Overrides the ``lookup_style`` given in the constructor.
        precision: Overrides the ``precision`` given in the constructor.
  
      Returns:
        Tensor of ``ids.shape + [embedding_dim]``.
  
      Raises:
        AttributeError: If ``lookup_style`` is not valid.
        ValueError: If ``ids`` is not an integer array.
      """
      # TODO(tomhennigan) Consider removing asarray here.
      ids = jnp.asarray(ids)
      if not jnp.issubdtype(ids.dtype, jnp.integer):
          raise ValueError(
              "hk.Embed's __call__ method must take an array of "
              "integer dtype but was called with an array of "
              f"{ids.dtype}"
          )
  
      lookup_style = lookup_style or self.lookup_style
      if isinstance(lookup_style, str):
          lookup_style = getattr(
              hk.EmbedLookupStyle, lookup_style.upper()
          )
  
      if lookup_style == hk.EmbedLookupStyle.ARRAY_INDEX:
          # If you don't wrap ids in a singleton tuple then JAX will try to unpack
          # it along the row dimension and treat each row as a separate index into
          # one of the dimensions of the array. The error only surfaces when
          # indexing with DeviceArray, while indexing with numpy.ndarray works fine.
          # See https://github.com/google/jax/issues/620 for more details.
          # Cast to a jnp array in case `ids` is a tracer (eg un a dynamic_unroll).
          return jnp.asarray(self.embeddings)[(ids,)]
  
      elif lookup_style == hk.EmbedLookupStyle.ONE_HOT:
          one_hot_ids = jax.nn.one_hot(ids, self.vocab_size)
          precision = self.precision if precision is None else precision
          return jnp.dot(
              one_hot_ids, self.embeddings, precision=precision
          )
  
      else:
          raise NotImplementedError(
              f"{lookup_style} is not supported by hk.Embed."
          )
  
file_path: dm-haiku/haiku/_src/embed.py
imports: from collections.abc import Sequence
  import enum
  from typing import Optional, Union
  from haiku._src import base
  from haiku._src import initializers
  from haiku._src import module
  import jax
  import jax.numpy
  import numpy
label_desc: Use image node_id 3 for calling the Embed obj's underlying member method code with example usage: obj.__call__(ids, lookup_style, precision) and returns: jnp, jnp