tensorflow/compiler/mlir/tensorflow/transforms/tpu_space_to_depth_pass.cc
// Handles cast op between the first convolution and the block argument.
Handles cast op between the first convolution and the block argument.
Handles cast op between the first convolution and the block argument .
LogicalResult HandleCast(TF::CastOp cast_op, ArrayRef<int64_t> new_shape) {
  auto cast_input = cast_op.x();
  auto transform_result_type =
      RankedTensorType::get(new_shape, getElementTypeOrSelf(cast_input));
  cast_input.setType(transform_result_type);
  auto block_arg = cast_input.dyn_cast<mlir::BlockArgument>();
  auto cast_op_input = dyn_cast_or_null<TF::CastOp>(cast_input.getDefiningOp());
  while (block_arg || cast_op_input) {
    if (block_arg) {
      HandleFuncOp(block_arg.getOwner()->getParentOp());
      block_arg = nullptr;
      cast_op_input = nullptr;
    } else {
      auto cast_input = cast_op_input.x();
      auto transform_result_type =
          RankedTensorType::get(new_shape, getElementTypeOrSelf(cast_input));
      cast_input.setType(transform_result_type);
      block_arg = cast_input.dyn_cast<mlir::BlockArgument>();
      cast_op_input = dyn_cast_or_null<TF::CastOp>(cast_input.getDefiningOp());
    }
  }
  return success();
}
LogicalResult HandleCast ( TF :: CastOp cast_op , ArrayRef < int64_t > new_shape ) { auto cast_input = cast_op . x ( ) ; auto transform_result_type = RankedTensorType :: get ( new_shape , getElementTypeOrSelf ( cast_input ) ) ; cast_input . setType ( transform_result_type ) ; auto block_arg = cast_input . dyn_cast < mlir :: BlockArgument > ( ) ; auto cast_op_input = dyn_cast_or_null < TF :: CastOp > ( cast_input . getDefiningOp ( ) ) ; while ( block_arg || cast_op_input ) { if ( block_arg ) { HandleFuncOp ( block_arg . getOwner ( ) -> getParentOp ( ) ) ; block_arg = nullptr ; cast_op_input = nullptr ; } else { auto cast_input = cast_op_input . x ( ) ; auto transform_result_type = RankedTensorType :: get ( new_shape , getElementTypeOrSelf ( cast_input ) ) ; cast_input . setType ( transform_result_type ) ; block_arg = cast_input . dyn_cast < mlir :: BlockArgument > ( ) ; cast_op_input = dyn_cast_or_null < TF :: CastOp > ( cast_input . getDefiningOp ( ) ) ; } } return success ( ) ; }
Handles cast op between the first convolution and the block argument.
Handles cast op between the first convolution and the block argument .
// Update input type. // Change on device function type/shape. // Update input type. // Update block arg and cast_op_input.