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: Attend over the embedding using a query array.
  
  Args:
    query: array with last dimension equal the feature depth `features` of the
      embedding.
  
  Returns:
    An array with final dim `num_embeddings` corresponding to the batched
    inner-product of the array of query vectors against each embedding.
    Commonly used for weight-sharing between embeddings and logit transform
    in NLP models.
docstr_tok: ["Attend","over","the","embedding","using","a","query","array",".","Args",":","query",":","array","with","last","dimension","equal","the","feature","depth","`","features","`","of","the","embedding",".","Returns",":","An","array","with","final","dim","`","num_embeddings","`","corresponding","to","the","batched","inner-product","of","the","array","of","query","vectors","against","each","embedding",".","Commonly","used","for","weight-sharing","between","embeddings","and","logit","transform","in","NLP","models","."]
code: def attend(self, query: Array) -> Array:
      """Attend over the embedding using a query array.
  
      Args:
        query: array with last dimension equal the feature depth `features` of the
          embedding.
  
      Returns:
        An array with final dim `num_embeddings` corresponding to the batched
        inner-product of the array of query vectors against each embedding.
        Commonly used for weight-sharing between embeddings and logit transform
        in NLP models.
      """
      dtype = (
          self.attend_dtype
          if self.attend_dtype is not None
          else self.dtype
      )
      return jnp.dot(query, jnp.asarray(self.embedding, dtype).T)
  
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 3 for calling the Embed obj's underlying member method code with example usage: obj.attend(query) and returns: jnp