for (Operation& inner_op :
       llvm::make_early_inc_range(replicate_op.GetBody())) {
    if (llvm::isa<tf_device::ReturnOp>(inner_op)) continue;
    // Skip hoisting if the inner op device attribute is a virtual device
    // defined by tf_device.replicate.
    if (DescendantUsesVirtualDevice(virtual_device_list, &inner_op)) continue;

    // We have to recurse, since we might have wrapped a side-effectful operator
    // in a tfl::CustomTfOp.
    if (op.walk([&](Operation *inner_op) {
            return op_has_side_effects_(inner_op) ? WalkResult::interrupt()
                                                  : WalkResult::advance();
          }).wasInterrupted()) {
      ops_with_side_effects.push_back(&op);
    }
  }