class_docstr: A simple lookup table that stores embeddings of a fixed size.
class_docstr_tok: ["A","simple","lookup","table","that","stores","embeddings","of","a","fixed","size","."]
docstr: **Arguments:**
  
  `Embedding` should be initialised with either:
  
  - `num_embeddings`: Size of embedding dictionary. Must be non-negative.
  - `embedding_size`: Size of each embedding vector. Must be non-negative.
  - `key`: A `jax.random.PRNGKey` used to provide randomness for initialisation
      of the embedding lookup table. (Keyword only argument.)
  
  Or:
  
  - `weight`: The embedding lookup table, of shape
      `(num_embeddings, embedding_size)`.
docstr_tok: ["*","*","Arguments",":","*","*","`","Embedding","`","should","be","initialised","with","either",":","-","`","num_embeddings","`",":","Size","of","embedding","dictionary",".","Must","be","non-negative",".","-","`","embedding_size","`",":","Size","of","each","embedding","vector",".","Must","be","non-negative",".","-","`","key","`",":","A","`","jax.random.PRNGKey","`","used","to","provide","randomness","for","initialisation","of","the","embedding","lookup","table",".","(","Keyword","only","argument",".",")","Or",":","-","`","weight","`",":","The","embedding","lookup","table",",","of","shape","`","(","num_embeddings",",","embedding_size",")","`","."]
code: def __init__(
      self,
      num_embeddings: Optional[int] = None,  # pyright: ignore
      embedding_size: Optional[int] = None,  # pyright: ignore
      weight: Optional[
          Float[Array, "num_embeddings embedding_size"]
      ] = None,
      *,
      key: Optional[PRNGKeyArray] = None,
  ):
      """**Arguments:**
  
      `Embedding` should be initialised with either:
  
      - `num_embeddings`: Size of embedding dictionary. Must be non-negative.
      - `embedding_size`: Size of each embedding vector. Must be non-negative.
      - `key`: A `jax.random.PRNGKey` used to provide randomness for initialisation
          of the embedding lookup table. (Keyword only argument.)
  
      Or:
  
      - `weight`: The embedding lookup table, of shape
          `(num_embeddings, embedding_size)`.
      """
      if weight is None:
          if (
              num_embeddings is None
              or embedding_size is None
              or key is None
          ):
              raise ValueError(
                  "Must provide `eqx.nn.Embedding(num_embeddings=..., "
                  "embedding_size=..., key=...)` if not providing the weight "
                  "directly."
              )
          if num_embeddings < 0:
              raise ValueError("num_embeddings must not be negative.")
          if embedding_size < 0:
              raise ValueError("embedding_size must not be negative.")
          self.weight = jrandom.normal(
              key, (num_embeddings, embedding_size)
          )
      else:
          if weight.ndim != 2:
              raise ValueError(
                  "weight must have shape (num_embeddings, embedding_size)."
              )
          if num_embeddings is None:
              num_embeddings: int = weight.shape[0]
          if embedding_size is None:
              embedding_size: int = weight.shape[1]
          if weight.shape != (num_embeddings, embedding_size):
              raise ValueError(
                  "weight must have shape (num_embeddings, embedding_size)."
              )
          self.weight = weight
      self.num_embeddings = num_embeddings
      self.embedding_size = embedding_size
  
file_path: equinox/equinox/nn/_embedding.py
imports: from typing import Optional
  import jax
  import jax.numpy
  import jax.random
  from jaxtyping import Array, ArrayLike, Float, Int, PRNGKeyArray
  from .._filters import is_array_like
  from .._module import field, Module
label_desc: Use image node_id 1 to create a new Embedding object from inherited base classes: Module with example: obj = Embedding(num_embeddings, embedding_size, weight)