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.
docstr: Embeds the inputs along the last dimension.
  
  Args:
    inputs: input data, all dimensions are considered batch dimensions.
      Values in the input array must be integers.
  
  Returns:
    Output which is embedded input data.  The output shape follows the input,
    with an additional `features` dimension appended.
docstr_tok: ["Embeds","the","inputs","along","the","last","dimension",".","Args",":","inputs",":","input","data",",","all","dimensions","are","considered","batch","dimensions",".","Values","in","the","input","array","must","be","integers",".","Returns",":","Output","which","is","embedded","input","data",".","The","output","shape","follows","the","input",",","with","an","additional","`","features","`","dimension","appended","."]
code: def __call__(self, inputs: Array) -> Array:
      """Embeds the inputs along the last dimension.
  
      Args:
        inputs: input data, all dimensions are considered batch dimensions.
          Values in the input array must be integers.
  
      Returns:
        Output which is embedded input data.  The output shape follows the input,
        with an additional `features` dimension appended.
      """
      if not jnp.issubdtype(inputs.dtype, jnp.integer):
          raise ValueError(
              "Input type must be an integer or unsigned integer."
          )
      # Use take because fancy indexing numpy arrays with JAX indices does not
      # work correctly.
      (embedding,) = dtypes.promote_dtype(
          self.embedding, dtype=self.dtype, inexact=False
      )
      if self.num_embeddings == 1:
          return jnp.where(
              jnp.broadcast_to(
                  inputs[..., None], inputs.shape + (self.features,)
              )
              == 0,
              embedding,
              jnp.nan,
          )
      return jnp.take(embedding, inputs, axis=0)
  
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 2 for calling the Embed obj's underlying member method code with example usage: obj.__call__(inputs) and returns: jnp, jnp