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: 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.
      """
      query, embedding = dtypes.promote_dtype(
          query, self.embedding, dtype=self.dtype
      )
      return jnp.dot(query, embedding.T)
  
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 3 for calling the Embed obj's underlying member method code with example usage: obj.attend(query) and returns: jnp