auto pack_op =
      rewriter.create<TF::PackOp>(loc, packed_type, /*values=*/matmuls, axis);

  // Reshape the rank-3 tensor into the correct output shape.
  const auto& result_batch_shape = bcast.output_batch_shape().dim_sizes();
  std::vector<int64_t> result_shape(result_batch_shape.begin(),
                                    result_batch_shape.end());
  result_shape.push_back(rows);
  result_shape.push_back(cols);