}
}  // namespace

mlir::LogicalResult WalkReachableFromTpuCluster(
    ModuleOp module, std::function<WalkResult(Operation*, tf_device::ClusterOp,
                                              std::optional<std::string>)>
                         callback) {
  return WalkReachableFromTpuCluster(true, module, callback);
}

mlir::LogicalResult WalkReachableFromTpuCluster(
    ModuleOp module,

// indicates failure.
// The host device is null when the tpu_cluster HasModelParallelism: The
// HasModelParallelism case is currently unsupported in combination with
// outside compilation.
mlir::LogicalResult WalkReachableFromTpuCluster(
    ModuleOp module, std::function<WalkResult(Operation*, tf_device::ClusterOp,
                                              std::optional<std::string>)>
                         callback);

            !isa<tf_device::LaunchOp>(op)) {
          op->removeAttr(kDeviceAttr);
        }
      }
      return WalkResult::advance();
    };

    if (failed(TFTPU::WalkReachableFromTpuCluster(getOperation(), traverse_op)))
      return signalPassFailure();
  }
};

}  // namespace

std::unique_ptr<OperationPass<ModuleOp>>
CreateTPUClusterCleanupAttributesPass() {