tensorflow/compiler/mlir/tensorflow/transforms/tpu_space_to_depth_pass.cc
// Gets input block argument id and number of users for the input recursively.
// Current supported ops between convolution input and the block arguments are
// PadOp and CastOp.
Gets input block argument id and number of users for the input recursively.
Current supported ops between convolution input and the block arguments are
PadOp and CastOp.
Gets input block argument id and number of users for the input recursively . Current supported ops between convolution input and the block arguments are PadOp and CastOp .
Optional<BlockArgumentInfo> GetInputBlockArgNum(Value input) {
  auto block_arg_num = GetBlockArgNum(input);
  if (block_arg_num.hasValue()) return block_arg_num;
  Value next_input = input;
  auto pad_op = dyn_cast_or_null<TF::PadOp>(next_input.getDefiningOp());
  auto cast_op = dyn_cast_or_null<TF::CastOp>(next_input.getDefiningOp());
  while (pad_op || cast_op) {
    if (pad_op) {
      auto block_arg_num = GetBlockArgNum(pad_op.input());
      if (block_arg_num.hasValue()) return block_arg_num;
      next_input = pad_op.input();
    } else {
      auto block_arg_num = GetBlockArgNum(cast_op.x());
      if (block_arg_num.hasValue()) return block_arg_num;
      next_input = cast_op.x();
    }
    pad_op = dyn_cast_or_null<TF::PadOp>(next_input.getDefiningOp());
    cast_op = dyn_cast_or_null<TF::CastOp>(next_input.getDefiningOp());
  }
  return None;
}
Optional < BlockArgumentInfo > GetInputBlockArgNum ( Value input ) { auto block_arg_num = GetBlockArgNum ( input ) ; if ( block_arg_num . hasValue ( ) ) return block_arg_num ; Value next_input = input ; auto pad_op = dyn_cast_or_null < TF :: PadOp > ( next_input . getDefiningOp ( ) ) ; auto cast_op = dyn_cast_or_null < TF :: CastOp > ( next_input . getDefiningOp ( ) ) ; while ( pad_op || cast_op ) { if ( pad_op ) { auto block_arg_num = GetBlockArgNum ( pad_op . input ( ) ) ; if ( block_arg_num . hasValue ( ) ) return block_arg_num ; next_input = pad_op . input ( ) ; } else { auto block_arg_num = GetBlockArgNum ( cast_op . x ( ) ) ; if ( block_arg_num . hasValue ( ) ) return block_arg_num ; next_input = cast_op . x ( ) ; } pad_op = dyn_cast_or_null < TF :: PadOp > ( next_input . getDefiningOp ( ) ) ; cast_op = dyn_cast_or_null < TF :: CastOp > ( next_input . getDefiningOp ( ) ) ; } return None ; }
Gets input block argument id and number of users for the input recursively.
Gets input block argument id and number of users for the input recursively .