tensorflow/lite/toco/graph_transformations/move_binary_operator_before_reshape.cc
// If a binary operator is doing a broadcast operation from a constant array,
// and the constant array shape is the tail of both the other input shape, and a
// subsequent reshape op's output shape, we can swap their order. Since we
// prefer to have reshape ops after mathematic ops, this can allow for the
// collapsing of some reshapes. The WaveNet model in particular benefits from
// this transformation.
//
// Note we are testing for one particular case of a broader set of possible
// binary-reshape op transformations. This transformation could be generalized.
If a binary operator is doing a broadcast operation from a constant array,
and the constant array shape is the tail of both the other input shape, and a
subsequent reshape op's output shape, we can swap their order. Since we
prefer to have reshape ops after mathematic ops, this can allow for the
collapsing of some reshapes. The WaveNet model in particular benefits from
this transformation.

Note we are testing for one particular case of a broader set of possible
binary-reshape op transformations. This transformation could be generalized.
If a binary operator is doing a broadcast operation from a constant array and the constant array shape is the tail of both the other input shape and a subsequent reshape op ' s output shape we can swap their order . Since we prefer to have reshape ops after mathematic ops this can allow for the collapsing of some reshapes . The WaveNet model in particular benefits from this transformation . Note we are testing for one particular case of a broader set of possible binary - reshape op transformations . This transformation could be generalized .
::tensorflow::Status MoveBinaryOperatorBeforeReshape::Run(Model* model,
                                                          std::size_t op_index,
                                                          bool* modified) {
  *modified = false;
  const auto binary_it = model->operators.begin() + op_index;
  Operator* binary_op = binary_it->get();
  if (binary_op->type != OperatorType::kAdd &&
      binary_op->type != OperatorType::kMul &&
      binary_op->type != OperatorType::kSub &&
      binary_op->type != OperatorType::kDiv &&
      binary_op->type != OperatorType::kFloorDiv &&
      binary_op->type != OperatorType::kFloorMod &&
      binary_op->type != OperatorType::kMinimum &&
      binary_op->type != OperatorType::kMaximum &&
      binary_op->type != OperatorType::kLess &&
      binary_op->type != OperatorType::kLessEqual &&
      binary_op->type != OperatorType::kGreater &&
      binary_op->type != OperatorType::kGreaterEqual) {
    return ::tensorflow::Status::OK();
  }
  CHECK_EQ(binary_op->inputs.size(), 2);
  const bool input_is_const[2] = {
      IsConstantParameterArray(*model, binary_op->inputs[0]),
      IsConstantParameterArray(*model, binary_op->inputs[1]),
  };
  if (!input_is_const[0] && !input_is_const[1]) {
    return ::tensorflow::Status::OK();
  }
  if (input_is_const[0] && input_is_const[1]) {
    return ::tensorflow::Status::OK();
  }
  const int constant_input_idx = input_is_const[0] ? 0 : 1;
  const int variable_input_idx = input_is_const[0] ? 1 : 0;
  CHECK(input_is_const[constant_input_idx]);
  CHECK(!input_is_const[variable_input_idx]);
  const auto& variable_input_array =
      model->GetArray(binary_op->inputs[variable_input_idx]);
  if (!variable_input_array.has_shape()) {
    AddMessageF(
        "Not moving %s because it's non-constant input shape is not resolved.",
        LogName(*binary_op));
    return ::tensorflow::Status::OK();
  }
  if (!IsTailOfShape(
          model->GetArray(binary_op->inputs[constant_input_idx]).shape(),
          model->GetArray(binary_op->inputs[variable_input_idx]).shape())) {
    return ::tensorflow::Status::OK();
  }
  auto reshape_it =
      FindOpWithOutput(*model, binary_op->inputs[variable_input_idx]);
  if (reshape_it == model->operators.end()) {
    AddMessageF("Not moving %s because it's variable input is not connected.",
                LogName(*binary_op));
    return ::tensorflow::Status::OK();
  }
  Operator* reshape_op = reshape_it->get();
  if (reshape_op->type != OperatorType::kReshape) {
    AddMessageF("Not moving %s because the preceding %s is not a reshape op",
                LogName(*binary_op), LogName(*reshape_op));
    return ::tensorflow::Status::OK();
  }
  const auto& reshape_input_array = model->GetArray(reshape_op->inputs[0]);
  if (!reshape_input_array.has_shape()) {
    AddMessageF(
        "Not moving %s because it's non-constant input shape is not resolved "
        "yet",
        LogName(*binary_op));
    return ::tensorflow::Status::OK();
  }
  if (!IsTailOfShape(
          model->GetArray(binary_op->inputs[constant_input_idx]).shape(),
          model->GetArray(reshape_op->outputs[0]).shape())) {
    return ::tensorflow::Status::OK();
  }
  for (const std::string& output_array : model->flags.output_arrays()) {
    if (binary_op->inputs[variable_input_idx] == output_array) {
      AddMessageF(
          "Not moving %s because the output of reshape op %s is an output op.",
          LogName(*binary_op), LogName(*reshape_op));
      return ::tensorflow::Status::OK();
    }
  }
  int count_ops_consuming_output =
      CountOpsWithInput(*model, binary_op->inputs[variable_input_idx]);
  DCHECK_GE(count_ops_consuming_output, 1);
  if (count_ops_consuming_output > 1) {
    AddMessageF(
        "Not moving %s because the output of reshape op %s is consumed by "
        "another op",
        LogName(*binary_op), LogName(*reshape_op));
    return ::tensorflow::Status::OK();
  }
  AddMessageF("Moving op %s before reshape op %s", LogName(*binary_op),
              LogName(*reshape_op));
  std::iter_swap(reshape_op->inputs.begin(),
                 binary_op->inputs.begin() + variable_input_idx);
  std::iter_swap(reshape_op->outputs.begin(), binary_op->outputs.begin());
  std::iter_swap(binary_it, reshape_it);
  model->GetArray(binary_op->outputs[0]).clear_shape();
  *modified = true;
  return ::tensorflow::Status::OK();
}
:: tensorflow :: Status MoveBinaryOperatorBeforeReshape :: Run ( Model * model , std :: size_t op_index , bool * modified ) { * modified = false ; const auto binary_it = model -> operators . begin ( ) + op_index ; Operator * binary_op = binary_it -> get ( ) ; if ( binary_op -> type != OperatorType :: kAdd && binary_op -> type != OperatorType :: kMul && binary_op -> type != OperatorType :: kSub && binary_op -> type != OperatorType :: kDiv && binary_op -> type != OperatorType :: kFloorDiv && binary_op -> type != OperatorType :: kFloorMod && binary_op -> type != OperatorType :: kMinimum && binary_op -> type != OperatorType :: kMaximum && binary_op -> type != OperatorType :: kLess && binary_op -> type != OperatorType :: kLessEqual && binary_op -> type != OperatorType :: kGreater && binary_op -> type != OperatorType :: kGreaterEqual ) { return :: tensorflow :: Status :: OK ( ) ; } CHECK_EQ ( binary_op -> inputs . size ( ) , 2 ) ; const bool input_is_const [ 2 ] = { IsConstantParameterArray ( * model , binary_op -> inputs [ 0 ] ) , IsConstantParameterArray ( * model , binary_op -> inputs [ 1 ] ) , } ; if ( ! input_is_const [ 0 ] && ! input_is_const [ 1 ] ) { return :: tensorflow :: Status :: OK ( ) ; } if ( input_is_const [ 0 ] && input_is_const [ 1 ] ) { return :: tensorflow :: Status :: OK ( ) ; } const int constant_input_idx = input_is_const [ 0 ] ? 0 : 1 ; const int variable_input_idx = input_is_const [ 0 ] ? 1 : 0 ; CHECK ( input_is_const [ constant_input_idx ] ) ; CHECK ( ! input_is_const [ variable_input_idx ] ) ; const auto & variable_input_array = model -> GetArray ( binary_op -> inputs [ variable_input_idx ] ) ; if ( ! variable_input_array . has_shape ( ) ) { AddMessageF ( " " , LogName ( * binary_op ) ) ; return :: tensorflow :: Status :: OK ( ) ; } if ( ! IsTailOfShape ( model -> GetArray ( binary_op -> inputs [ constant_input_idx ] ) . shape ( ) , model -> GetArray ( binary_op -> inputs [ variable_input_idx ] ) . shape ( ) ) ) { return :: tensorflow :: Status :: OK ( ) ; } auto reshape_it = FindOpWithOutput ( * model , binary_op -> inputs [ variable_input_idx ] ) ; if ( reshape_it == model -> operators . end ( ) ) { AddMessageF ( " " , LogName ( * binary_op ) ) ; return :: tensorflow :: Status :: OK ( ) ; } Operator * reshape_op = reshape_it -> get ( ) ; if ( reshape_op -> type != OperatorType :: kReshape ) { AddMessageF ( " " , LogName ( * binary_op ) , LogName ( * reshape_op ) ) ; return :: tensorflow :: Status :: OK ( ) ; } const auto & reshape_input_array = model -> GetArray ( reshape_op -> inputs [ 0 ] ) ; if ( ! reshape_input_array . has_shape ( ) ) { AddMessageF ( " " " " , LogName ( * binary_op ) ) ; return :: tensorflow :: Status :: OK ( ) ; } if ( ! IsTailOfShape ( model -> GetArray ( binary_op -> inputs [ constant_input_idx ] ) . shape ( ) , model -> GetArray ( reshape_op -> outputs [ 0 ] ) . shape ( ) ) ) { return :: tensorflow :: Status :: OK ( ) ; } for ( const std :: string & output_array : model -> flags . output_arrays ( ) ) { if ( binary_op -> inputs [ variable_input_idx ] == output_array ) { AddMessageF ( " " , LogName ( * binary_op ) , LogName ( * reshape_op ) ) ; return :: tensorflow :: Status :: OK ( ) ; } } int count_ops_consuming_output = CountOpsWithInput ( * model , binary_op -> inputs [ variable_input_idx ] ) ; DCHECK_GE ( count_ops_consuming_output , 1 ) ; if ( count_ops_consuming_output > 1 ) { AddMessageF ( " " " " , LogName ( * binary_op ) , LogName ( * reshape_op ) ) ; return :: tensorflow :: Status :: OK ( ) ; } AddMessageF ( " " , LogName ( * binary_op ) , LogName ( * reshape_op ) ) ; std :: iter_swap ( reshape_op -> inputs . begin ( ) , binary_op -> inputs . begin ( ) + variable_input_idx ) ; std :: iter_swap ( reshape_op -> outputs . begin ( ) , binary_op -> outputs . begin ( ) ) ; std :: iter_swap ( binary_it , reshape_it ) ; model -> GetArray ( binary_op -> outputs [ 0 ] ) . clear_shape ( ) ; * modified = true ; return :: tensorflow :: Status :: OK ( ) ; }
If a binary operator is doing a broadcast operation from a constant array,
and the constant array shape is the tail of both the other input shape, and a
subsequent reshape op's output shape, we can swap their order.
If a binary operator is doing a broadcast operation from a constant array and the constant array shape is the tail of both the other input shape and a subsequent reshape op ' s output shape we can swap their order .
// BINARY OP INPUT CHECKS // To limit our scope, we require one constant input. Though there's no // reason this transformation wouldn't work with all variable inputs. // Both inputs are constants. Leave this for constants propagation. // Constant array shape must be the latter part of the variable shape. // RESHAPE OP CHECKS // Constant array shape must be the latter part of the binary op output // shape. // EXTRA CHECKS ON CONNECTING ARRAY // SWAP ORDER OF BINARY AND RESHAPE OPS // Swap op input and outputs // Swap operator ordering // Clear binary output shape so it will be re-propagated