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):
      """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.
      """
      return jnp.dot(query, self.table.T)
  
file_path: flax/flax/core/nn/linear.py
imports: from collections.abc import Iterable
  import jax.numpy
  import numpy
  from jax import lax
  from flax import struct
  from flax.core import Scope
  from flax.linen import initializers
label_desc: Use image node_id 2 for calling the Embedding obj's underlying member method code with example usage: obj.attend(query) and returns: jnp