Function ReshapeToRank3

tensorflow/core/kernels/einsum_op.cc:510–515

Reshapes a Tensor of shape [b0,b1...bk,N,M] to [prod(b0,b1...bk),N,M].

// Reshapes a Tensor of shape [b0,b1...bk,N,M] to [prod(b0,b1...bk),N,M].
Status ReshapeToRank3(const Tensor& input, int batch_size, Tensor* output) {
  const int rank = input.dims();
  TensorShape output_shape = {batch_size, input.dim_size(rank - 2),
                              input.dim_size(rank - 1)};
  return CopyFrom(input, output_shape, output);
}

// Contracts the inputs along the last axis. (or the second last if the
// corresponding value of swap_free_and_contract is true). The batch dimensions