code: def compute_q(self, query_antecedent):
      
      ret = mtf.einsum(
          [query_antecedent, self.wq], reduced_dims=[self.query_input_dim])
      if self.combine_dims:
        ret = mtf.replace_dimensions(ret, ret.shape.dims[-1], self.q_dims)
      return ret
docs: Compute query Tensor q.
  
      Args:
        query_antecedent: a Tensor with dimensions
           {query_input_dim} + other_dims
      Returns:
        a Tensor with dimensions
           query_heads_dims + {key_dim} + other_dims