code: def compute_kv(self, memory_antecedent):
      
      if not self.shared_kv:
        raise ValueError("compute_kv can only be called with shared_kv")
      ret = mtf.einsum(
          [memory_antecedent, self.wkv], reduced_dims=[self.memory_input_dim])
      if self.combine_dims:
        ret = mtf.replace_dimensions(ret, ret.shape.dims[-1], self.k_dims)
      return ret
docs: Compute key/value Tensor kv.
  
      Args:
        memory_antecedent: a Tensor with dimensions
          {memory_input_dim} + other_dims
      Returns:
        a Tensor with dimensions
          memory_heads_dims + {key_dim} + other_dims