return rewriter->create<TF::ReshapeOp>(loc, out_tensor_type,
                                         /*tensor=*/value,
                                         /*shape=*/segprod->getResults()[0]);
}

struct EinsumDimensionNumbers {
  // Each field contains the list of dimensions appearing only in the specifed
  // arguments of the einsum op with natural ordering. For example `rhs_out`