code: def compute_output(self, o, output_shape=None):
      
      if self.combine_dims:
        o = mtf.transpose(o, o.shape - self.o_dims + self.o_dims)
        o = mtf.replace_dimensions(o, self.o_dims, self.wo.shape.dims[0])
        reduced_dims = [self.wo.shape.dims[0]]
      else:
        reduced_dims = self.o_dims
      return mtf.einsum(
          [o, self.wo], output_shape=output_shape, reduced_dims=reduced_dims)
docs: Compute output of multihead attention.
  
      Args:
        o: a Tensor with dimensions
           query_heads_dims + {value_dim} + other_dims
        output_shape: an optional Shape
      Returns:
        a Tensor with shape:
           {output_dim} + other_dims