Method Run

tensorflow/lite/toco/graph_transformations/unroll_batch_matmul.cc:133–259

Unrolls a BatchMatMul on the batch dimension. We need to slice each batch out of the inputs, matmul them individually, then stack them all back together at the end.

// We need to slice each batch out of the inputs, matmul them individually, then
// stack them all back together at the end.
::tensorflow::Status UnrollBatchMatMul::Run(Model* model, std::size_t op_index,
                                            bool* modified) {
  *modified = false;
  auto batch_op_it = model->operators.begin() + op_index;
  if (batch_op_it->get()->type != OperatorType::kBatchMatMul) {
    return ::tensorflow::Status::OK();
  }
  const auto* batch_op =
      static_cast<const BatchMatMulOperator*>(batch_op_it->get());
  auto& tail_it = batch_op_it;

  string input_lhs = batch_op->inputs[0];
  string input_rhs = batch_op->inputs[1];
  const auto& input_lhs_array = model->GetArray(input_lhs);
  const auto& input_rhs_array = model->GetArray(input_rhs);
  if (!input_lhs_array.has_shape() || !input_rhs_array.has_shape())
    return ::tensorflow::Status::OK();

  // Transpose LHS input if necessary.
  if (batch_op->adj_x) {
    TransposeOperator* transpose_op = TransposeInput(input_lhs, model);
    tail_it = model->operators.emplace(tail_it, transpose_op) + 1;
    input_lhs = transpose_op->outputs[0];
  }
  const auto& input_array_a = model->GetArray(input_lhs);

  // Transpose RHS input if necessary.
  if (batch_op->adj_y) {
    TransposeOperator* transpose_op = TransposeInput(input_rhs, model);
    tail_it = model->operators.emplace(tail_it, transpose_op) + 1;
    input_rhs = transpose_op->outputs[0];
  }
  const auto& input_array_b = model->GetArray(input_rhs);

  // Ensure that input ranks are at least 2 and batch shapes are broadcastable.
  const int dims_a = input_array_a.shape().dimensions_count();
  const int dims_b = input_array_b.shape().dimensions_count();
  CHECK_GE(dims_a, 2) << "First input must have rank >= 2";
  CHECK_GE(dims_b, 2) << "Second input must have rank >= 2";

  ::tensorflow::MatMulBCast bcast(
      ToInlinedVector(input_array_a.shape().dims()),
      ToInlinedVector(input_array_b.shape().dims()));
  CHECK(bcast.IsValid()) << "Input batch dimensions must be broadcastable";

  CHECK_EQ(input_array_a.shape().dims(dims_a - 1),
           input_array_b.shape().dims(dims_b - 2))
      << "Input dimensions must be compatible for multipication. shape a = ["
      << absl::StrJoin(input_array_a.shape().dims(), ", ") << "], shape b = ["
      << absl::StrJoin(input_array_b.shape().dims(), ", ") << "]";

  if (dims_a == 2 && dims_b == 2) {
    // This is really just a MatMul.
    AddMessageF("Replacing non-batch BatchMatMul %s by a MatMul operator",
                LogName(*batch_op));
    auto* matmul_op = new TensorFlowMatMulOperator;
    matmul_op->inputs = {input_lhs, input_rhs};
    matmul_op->outputs = batch_op->outputs;