class_docstr: Embedding Module.
  
  A parameterized function from integers [0, ``num_embeddings``) to
  ``features``-dimensional vectors. This ``Module`` will create an ``embedding``
  matrix with shape ``(num_embeddings, features)``. When calling this layer,
  the input values will be used to 0-index into the ``embedding`` matrix.
  Indexing on a value greater than or equal to ``num_embeddings`` will result
  in ``nan`` values. When ``num_embeddings`` equals to 1, it will
  broadcast the ``embedding`` matrix to input shape with ``features``
  dimension appended.
  
  Attributes:
    num_embeddings: number of embeddings / vocab size.
    features: number of feature dimensions for each embedding.
    dtype: the dtype of the embedding vectors (default: same as embedding).
    param_dtype: the dtype passed to parameter initializers (default: float32).
    embedding_init: embedding initializer.
code: def __init__(
      self,
      num_embeddings: int,
      features: int,
      *,
      dtype: tp.Optional[Dtype] = None,
      param_dtype: Dtype = jnp.float32,
      embedding_init: tp.Callable[
          [KeyArray, Shape, Dtype], Array
      ] = default_embed_init,
      rngs: rnglib.Rngs,
  ):
      self.embedding = nnx.Param(
          embedding_init(
              rngs.params(), (num_embeddings, features), param_dtype
          )
      )
  
      self.num_embeddings = num_embeddings
      self.features = features
      self.dtype = dtype or self.embedding.dtype
      self.param_dtype = param_dtype
      self.embedding_init = embedding_init
  
imports: from __future__ import annotations
  import typing
  from types import MappingProxyType
  import jax
  import jax.numpy
  import numpy
  from jax import lax
  from flax.experimental import nnx
  from flax.experimental.nnx.nnx import rnglib, variables
  from flax.experimental.nnx.nnx.module import Module
  from flax.experimental.nnx.nnx.nn import dtypes, initializers
label_desc: Use image node_id 1 to create a new Embed object from inherited base classes: Module with example: obj = Embed(num_embeddings, features)