class_docstr: Module for embedding tokens in a low-dimensional space.
class_docstr_tok: ["Module","for","embedding","tokens","in","a","low-dimensional","space","."]
docstr: Constructs an Embed module.
  
  Args:
    vocab_size: The number of unique tokens to embed. If not provided, an
      existing vocabulary matrix from which ``vocab_size`` can be inferred
      must be provided as ``embedding_matrix``.
    embed_dim: Number of dimensions to assign to each embedding. If an
      existing vocabulary matrix initializes the module, this should not be
      provided as it will be inferred.
    embedding_matrix: A matrix-like object equivalent in size to
      ``[vocab_size, embed_dim]``. If given, it is used as the initial value
      for the embedding matrix and neither ``vocab_size`` or ``embed_dim``
      need be given. If they are given, their values are checked to be
      consistent with the dimensions of ``embedding_matrix``.
    w_init: An initializer for the embeddings matrix. As a default, embeddings
      are initialized via a truncated normal distribution.
    lookup_style: One of the enum values of :class:`EmbedLookupStyle`
      determining how to access the value of the embeddings given an ID.
      Regardless the input should be a dense array of integer values
      representing ids. This setting changes how internally this module maps
      those ids to embeddings. The result is the same, but the speed and
      memory tradeoffs are different. It defaults to using NumPy-style array
      indexing. This value is only the default for the module, and at any
      given invocation can be overridden in :meth:`__call__`.
    name: Optional name for this module.
    precision: Only used when lookup_style is ONE_HOT. The precision to use
      for the dot-product between the one-hot-encoded inputs and the embedding
      vectors. It is possible to attain a ~2x speedup on TPU using
      `jax.lax.Precision.DEFAULT` at the cost of a slightly lower precision.
  
  Raises:
    ValueError: If none of ``embed_dim``, ``embedding_matrix`` and
      ``vocab_size`` are supplied, or if ``embedding_matrix`` is supplied
      and ``embed_dim`` or ``vocab_size`` is not consistent with the
      supplied matrix.
docstr_tok: ["Constructs","an","Embed","module",".","Args",":","vocab_size",":","The","number","of","unique","tokens","to","embed",".","If","not","provided",",","an","existing","vocabulary","matrix","from","which","``","vocab_size","``","can","be","inferred","must","be","provided","as","``","embedding_matrix","``",".","embed_dim",":","Number","of","dimensions","to","assign","to","each","embedding",".","If","an","existing","vocabulary","matrix","initializes","the","module",",","this","should","not","be","provided","as","it","will","be","inferred",".","embedding_matrix",":","A","matrix-like","object","equivalent","in","size","to","``","[","vocab_size",",","embed_dim","]","``",".","If","given",",","it","is","used","as","the","initial","value","for","the","embedding","matrix","and","neither","``","vocab_size","``","or","``","embed_dim","``","need","be","given",".","If","they","are","given",",","their","values","are","checked","to","be","consistent","with","the","dimensions","of","``","embedding_matrix","``",".","w_init",":","An","initializer","for","the","embeddings","matrix",".","As","a","default",",","embeddings","are","initialized","via","a","truncated","normal","distribution",".","lookup_style",":","One","of","the","enum","values","of",":","class",":","`","EmbedLookupStyle","`","determining","how","to","access","the","value","of","the","embeddings","given","an","ID",".","Regardless","the","input","should","be","a","dense","array","of","integer","values","representing","ids",".","This","setting","changes","how","internally","this","module","maps","those","ids","to","embeddings",".","The","result","is","the","same",",","but","the","speed","and","memory","tradeoffs","are","different",".","It","defaults","to","using","NumPy-style","array","indexing",".","This","value","is","only","the","default","for","the","module",",","and","at","any","given","invocation","can","be","overridden","in",":","meth",":","`","__call__","`",".","name",":","Optional","name","for","this","module",".","precision",":","Only","used","when","lookup_style","is","ONE_HOT",".","The","precision","to","use","for","the","dot-product","between","the","one-hot-encoded","inputs","and","the","embedding","vectors",".","It","is","possible","to","attain","a","~2x","speedup","on","TPU","using","`","jax.lax.Precision.DEFAULT","`","at","the","cost","of","a","slightly","lower","precision",".","Raises",":","ValueError",":","If","none","of","``","embed_dim","``",",","``","embedding_matrix","``","and","``","vocab_size","``","are","supplied",",","or","if","``","embedding_matrix","``","is","supplied","and","``","embed_dim","``","or","``","vocab_size","``","is","not","consistent","with","the","supplied","matrix","."]
code: def __init__(
      self,
      vocab_size: Optional[int] = None,
      embed_dim: Optional[int] = None,
      embedding_matrix: Optional[Union[np.ndarray, jax.Array]] = None,
      w_init: Optional[hk.initializers.Initializer] = None,
      lookup_style: Union[str, EmbedLookupStyle] = "ARRAY_INDEX",
      name: Optional[str] = None,
      precision: jax.lax.Precision = jax.lax.Precision.HIGHEST,
  ):
      """Constructs an Embed module.
  
      Args:
        vocab_size: The number of unique tokens to embed. If not provided, an
          existing vocabulary matrix from which ``vocab_size`` can be inferred
          must be provided as ``embedding_matrix``.
        embed_dim: Number of dimensions to assign to each embedding. If an
          existing vocabulary matrix initializes the module, this should not be
          provided as it will be inferred.
        embedding_matrix: A matrix-like object equivalent in size to
          ``[vocab_size, embed_dim]``. If given, it is used as the initial value
          for the embedding matrix and neither ``vocab_size`` or ``embed_dim``
          need be given. If they are given, their values are checked to be
          consistent with the dimensions of ``embedding_matrix``.
        w_init: An initializer for the embeddings matrix. As a default, embeddings
          are initialized via a truncated normal distribution.
        lookup_style: One of the enum values of :class:`EmbedLookupStyle`
          determining how to access the value of the embeddings given an ID.
          Regardless the input should be a dense array of integer values
          representing ids. This setting changes how internally this module maps
          those ids to embeddings. The result is the same, but the speed and
          memory tradeoffs are different. It defaults to using NumPy-style array
          indexing. This value is only the default for the module, and at any
          given invocation can be overridden in :meth:`__call__`.
        name: Optional name for this module.
        precision: Only used when lookup_style is ONE_HOT. The precision to use
          for the dot-product between the one-hot-encoded inputs and the embedding
          vectors. It is possible to attain a ~2x speedup on TPU using
          `jax.lax.Precision.DEFAULT` at the cost of a slightly lower precision.
  
      Raises:
        ValueError: If none of ``embed_dim``, ``embedding_matrix`` and
          ``vocab_size`` are supplied, or if ``embedding_matrix`` is supplied
          and ``embed_dim`` or ``vocab_size`` is not consistent with the
          supplied matrix.
      """
      super().__init__(name=name)
      if embedding_matrix is None and not (vocab_size and embed_dim):
          raise ValueError(
              "hk.Embed must be supplied either with an initial `embedding_matrix` "
              "or with `embed_dim` and `vocab_size`."
          )
  
      if embedding_matrix is not None:
          embedding_matrix = jnp.asarray(embedding_matrix)
          if vocab_size and embedding_matrix.shape[0] != vocab_size:
              raise ValueError(
                  "An `embedding_matrix` was supplied but the `vocab_size` of "
                  f"{vocab_size} was not consistent with its shape "
                  f"{embedding_matrix.shape}."
              )
          if embed_dim and embedding_matrix.shape[1] != embed_dim:
              raise ValueError(
                  "An `embedding_matrix` was supplied but the `embed_dim` of "
                  f"{embed_dim} was not consistent with its shape "
                  f"{embedding_matrix.shape}."
              )
          w_init = lambda _, __: embedding_matrix
          vocab_size = embedding_matrix.shape[0]
          embed_dim = embedding_matrix.shape[1]
  
      self.vocab_size = vocab_size
      self.embed_dim = embed_dim
      self.lookup_style = lookup_style
      self.precision = precision
      self.w_init = w_init or hk.initializers.TruncatedNormal()
  
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 1 to create a new Embed object from inherited base classes: hk with example: obj = Embed(vocab_size, embed_dim, embedding_matrix, w_init, lookup_style, name, precision)