_initializers = _KerasLazyLoader(
    globals(), submodule="initializers", name="initializers", mode="v2")
setattr(_current_module, "losses", _losses)
setattr(_current_module, "metrics", _metrics)
setattr(_current_module, "optimizers", _optimizers)

  auto cluster = builder.create<mlir::tf_device::ClusterOp>(
      mlir::UnknownLoc::get(&context), result_types);
  cluster->setAttr(kNumCoresPerReplicaAttr,
                   builder.getIntegerAttr(builder.getIntegerType(64), 1));
  cluster->setAttr(kTopologyAttr, builder.getStringAttr(""));
  cluster->setAttr(kDeviceAssignmentAttr, builder.getArrayAttr({}));

  EXPECT_FALSE(HasModelParallelism(cluster));
}

  added_call_op->setAttr(kInferenceType, inference_type);
  added_func_op->setAttr(kDevice, device);
  added_func_op->setAttr(kInferenceType, inference_type);

  std::string function_name = absl::StrCat(interface_name.getValue().str(), "_",
                                           device.getValue().str(), "_",
                                           inference_type.getValue().str());

  mlir::MLIRContext context;
  mlir::OwningOpRef<mlir::ModuleOp> module_ref =
      mlir::ModuleOp::create(mlir::UnknownLoc::get(&context));
  mlir::Builder builder(*module_ref);
  (*module_ref)->setAttr("tf.devices", builder.getBoolAttr(false));

  mlir::TF::RuntimeDevices devices;
  EXPECT_TRUE(mlir::failed(GetDevicesFromOp(*module_ref, &devices)));
}

    if (op->getDialect() != tf_dialect) return WalkResult::advance();

    if (parallel_group_attr) {
      op->setAttr(TF::kParallelExecAnnotation, parallel_group_attr);
    }
    auto device_attr = op->getAttr(kDeviceAttr);
    if (!device_attr) {
      op->setAttr(kDeviceAttr, launch.getDeviceAttr());
      return WalkResult::advance();
    }

_initializers = _KerasLazyLoader(
    globals(), submodule="initializers", name="initializers")
setattr(_current_module, "losses", _losses)
setattr(_current_module, "metrics", _metrics)
setattr(_current_module, "optimizers", _optimizers)
setattr(_current_module, "initializers", _initializers)


# Do an eager load for Keras' code so that any function/method that needs to

  // function name can be changed by those passes.
  call_op->setAttr(TF::kStablehloEntryFunctionAttrName,
                   FlatSymbolRefAttr::get(builder.getStringAttr(func_name)));
  // Store the custom attribute to restore the function name when loading it
  // back in the post calibration stage. As mentioned above, the above entry
  // function attribute is not reliable.
  call_op->setAttr(kOriginalStablehloEntryFunctionAttrName,

                                          dq.getArg());
      dq.getResult().replaceAllUsesWith(dcast);
      if (auto extra_attr = op->getAttr(mlir::quant::kVolatileOpAttrName)) {
        dcast->setAttr(mlir::quant::kVolatileOpAttrName, extra_attr);
      }
      dq.erase();
    } else if (auto q = llvm::dyn_cast<quantfork::QuantizeCastOp>(op)) {
      auto out_type = q.getResult().getType();

# Hook external TensorFlow modules.
_current_module = _sys.modules[__name__]

# Lazy load Keras v1
_tf_uses_legacy_keras = (
    _os.environ.get("TF_USE_LEGACY_KERAS", None) in ("true", "True", "1"))
setattr(_current_module, "keras", _KerasLazyLoader(globals(), mode="v1"))
if _tf_uses_legacy_keras:
  _module_dir = _module_util.get_parent_dir_for_name("tf_keras.api._v1.keras")
else:

};

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) {
        op->setAttr("_wraps_compilation", UnitAttr::get(func->getContext()));
        op = op->getParentOp();
      }
    }
  });
}