class_docstr: An embedding layer with named axes for :meth:`jax.experimental.pjit.pjit`.
  
  .. warning:: This class is hightly EXPERIMENTAL and the API is likely to
      change. For regular (non-pjit) use, please use
      :class:`flax.linen.linear.Embed`.
  
  Attributes:
    num_embeddings: number of embeddings.
    features: number of feature dimensions for each embedding.
    dtype: the dtype of the embedding vectors (default: float32).
    param_dtype: the dtype passed to parameter initializers (default: float32).
    embedding_init: embedding initializer.
    one_hot: performs the gather with a one-hot contraction rather than a true
      gather. This is currently needed for SPMD partitioning.
docstr: Embeds the inputs along the last dimension.
  
  Args:
    inputs: input data, all dimensions are considered batch dimensions.
  
  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",".","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.
  
      Returns:
        Output which is embedded input data.  The output shape follows the input,
        with an additional `features` dimension appended.
      """
      if self.cast_input_dtype:
          inputs = inputs.astype(self.cast_input_dtype)
      if not jnp.issubdtype(inputs.dtype, jnp.integer):
          raise ValueError(
              "Input type must be an integer or unsigned integer."
          )
      if self.one_hot:
          iota = lax.iota(jnp.int32, self.num_embeddings)
          one_hot = jnp.array(
              inputs[..., jnp.newaxis] == iota, dtype=self.dtype
          )
          output = jnp.dot(
              one_hot, jnp.asarray(self.embedding, self.dtype)
          )
      else:
          output = jnp.asarray(self.embedding, self.dtype)[inputs]
          output = with_sharding_constraint(
              output, ("batch", "length", "embed")
          )
      return output
  
imports: import dataclasses
  from typing import Any, Callable, Iterable, Optional, Tuple, Union
  import jax.numpy
  from jax import lax
  from flax import linen
  from flax.linen import initializers
  from flax.linen.linear import DotGeneralT, PrecisionLike
  from flax.linen.partitioning import param_with_axes, with_sharding_constraint
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: output