namespace stablehlo::quantization {

void DisableDebugging(mlir::ModuleOp module_op) {
  module_op.walk(
      [](mlir::TF::DumpTensorOp dump_op) { dump_op.setEnabled(false); });
}

void ChangeToQuantizedFilename(mlir::ModuleOp module_op) {
  module_op.walk([](mlir::TF::DumpTensorOp dump_op) {
    dump_op.setFileName("quantized_tensor_data.pb");
  });
}

struct MoveTpuCompileToFrontPass
    : public impl::MoveTpuCompileToFrontPassBase<MoveTpuCompileToFrontPass> {
  void runOnOperation() override;
};

void MarkCompilationOps(Operation* func) {
  func->walk([&](Operation* op) {
    if (llvm::isa<TF::_TPUCompileMlirOp>(op)) {
      op->setAttr("_is_compilation", UnitAttr::get(func->getContext()));
      op = op->getParentOp();
      while (op && op != func) {

  OpBuilder globalBuilder(module.getBodyRegion());

  module.walk(
      [&](TF::ReadVariableOp op) { MaybeCreateVarHandleForOp(op, solver); });
  module.walk(
      [&](TF::AssignVariableOp op) { MaybeCreateVarHandleForOp(op, solver); });
  module.walk(
      [&](TF::IteratorGetNextOp op) { MaybeCreateVarHandleForOp(op, solver); });
}

}  // namespace

// from filepath.Walk:
//   - List of files/dirs is not sorted for performance reasons
//   - callback walkFunc is invoked on root directory after visiting children dirs and files
func walkDeep(root string, walkFunc filepath.WalkFunc) error {
	info, err := os.Lstat(root)
	if err != nil {
		return walkFunc(root, nil, err)
	}
	return walk(root, info, walkFunc)
}

import org.hamcrest.CoreMatchers.equalTo
import org.hamcrest.MatcherAssert.assertThat

import org.junit.Test


class WalkReproduciblyTest : FolderBasedTest() {

    @Test
    fun `walks directory tree top-down yielding sorted files`() {
        withFolders {
            "root" {
                withFile("root-f2")
                "b" {
                    withFile("b-f1")
                }

        // We cannot inject internal services in to the transform directly, but we can create them via object factory
        ClasspathWalker walker = internalServices.get().getClasspathWalker();
        walker.visit(artifact, entry -> {
            if (entry.getName().endsWith(".class") && !isInUnsupportedMrJarVersionedDirectory(entry)) {

// results.
void RemoveIdentityFromRegion(Region& region) {
  region.walk([](Operation* op) {
    if (isa<TF::IdentityOp, TF::IdentityNOp>(op)) {
      op->replaceAllUsesWith(op->getOperands());
      op->erase();
    }
  });
}

void TPUIdentityPruning::runOnOperation() {
  SmallVector<tf_device::ClusterOp, 4> clusters;
  getOperation().walk(
      [&](tf_device::ClusterOp cluster) { clusters.push_back(cluster); });

      CallGraphNode* node = call_graph.resolveCallable(call, symbol_table);
      pending_call_nodes.emplace_back(node, tpu_cluster);
    }
  };

  // Traverse ops in each TPU cluster.
  auto result = module.walk([&](tf_device::ClusterOp tpu_cluster) {
    std::optional<std::string> host_device;
    if (pass_host_device) {
      std::string host_device_value;
      if (failed(tensorflow::GetHostDeviceOutsideComputation(

}

bool HasChildLaunchDeviceOp(Operation& op) {
  auto cluster_op = dyn_cast<tf_device::ClusterOp>(op);
  if (cluster_op == nullptr) return false;

  auto walk_result = cluster_op->walk([&](Operation* op) {
    if (IsLaunchOp(*op)) return WalkResult::interrupt();

    return WalkResult::advance();
  });

  return walk_result.wasInterrupted();
}

  module->walk([&](mlir::tf_device::ClusterFuncOp cluster_func) {
    cluster_func_ops.push_back(cluster_func);
  });
  EXPECT_EQ(cluster_func_ops.size(), 1);

  EXPECT_TRUE(mlir::succeeded(tensorflow::EraseClusterFuncs(cluster_func_ops)));

  llvm::SmallVector<mlir::tf_device::ClusterFuncOp, 4> new_cluster_func_ops;
  module->walk([&](mlir::tf_device::ClusterFuncOp cluster_func) {