class_docstr: A set of parameters used for (multihead) attention.
docstr: Compute key Tensor k.
  
  Args:
    memory_antecedent: a Tensor with dimensions
      {memory_input_dim} + other_dims
  Returns:
    a Tensor with dimensions
      memory_heads_dims + {key_dim} + other_dims
docstr_tok: ["Compute","key","Tensor","k.","Args",":","memory_antecedent",":","a","Tensor","with","dimensions","{","memory_input_dim","}","+","other_dims","Returns",":","a","Tensor","with","dimensions","memory_heads_dims","+","{","key_dim","}","+","other_dims"]
code: def compute_k(self, memory_antecedent):
      """Compute key Tensor k.
  
      Args:
        memory_antecedent: a Tensor with dimensions
          {memory_input_dim} + other_dims
      Returns:
        a Tensor with dimensions
          memory_heads_dims + {key_dim} + other_dims
      """
      if self.shared_kv:
          raise ValueError("compute_k cannot be called with shared_kv")
      ret = mtf.layers.us_einsum(
          [memory_antecedent, self.wk],
          reduced_dims=[self.memory_input_dim],
      )
      if self.combine_dims:
          ret = mtf.replace_dimensions(
              ret, ret.shape.dims[-1], self.k_dims
          )
      return ret
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import mesh_tensorflow
  import tensorflow.compat.v1
label_desc: Use image node_id 5 for calling the AttentionParams obj's underlying member method code with example usage: obj.compute_k(memory_antecedent) and returns: ret