func.return %0 : tensor<i32>
}
)mlir";
  mlir::MLIRContext context;
  context.loadDialect<mlir::func::FuncDialect, mlir::TF::TensorFlowDialect>();
  mlir::OwningOpRef<mlir::ModuleOp> module =
      mlir::parseSourceString<mlir::ModuleOp>(code, &context);
  ASSERT_TRUE(module);
  EXPECT_TRUE(IsSupportedByReplicatedBridge(*module));
}

TEST(HasTPUPartitionedCallOpInModule, HasTPUPartitionedCallModule) {

        absl::StrCat("Failed to open input file: ", error_message));
  }

  if (input_mlir) {
    source_mgr->AddNewSourceBuffer(std::move(file), llvm::SMLoc());
    return OwningOpRef<ModuleOp>(
        mlir::parseSourceFile<mlir::ModuleOp>(*source_mgr, context));
  }

  // Register extra TF ops passed as OpDef.
  auto extra_opdefs_status = RegisterExtraTfOpDefs(extra_tf_opdefs);

// vhlo.op %1 : vhlo.tensor<...>
//   ==>
// vhlo.op %1 : tensor<...>
//
// TODO: There's likely a way to make MLIR manage the unrealized cast
// conversions using a specific rewriter.
LogicalResult ApplyTypeConverter(ModuleOp op, TypeConverter &converter) {
  IRRewriter rewriter(op->getContext());

  op->walk([&](Operation *op) {
    if (op->getDialect()->getNamespace() != "vhlo") return;

    // Convert operands

    // and use it for `shape_inference_graph` attribute on XlaHostCompute.
    func::FuncOp cloned_func;
    SymbolTable manager(op->getParentOfType<ModuleOp>());
    StringRef host_module = op.getHostMlirModule();
    if (!host_module.empty()) {
      mlir::OwningOpRef<mlir::ModuleOp> module_for_func;

      func::FuncOp func = op.GetHostFunc(&module_for_func);

      OpBuilder::InsertionGuard guard(rewriter);

    Pass<"tf-resource-op-lifting-for-main-function", "ModuleOp"> {
  let summary = "Lifting resource operations out of control flow statements "
    "for the main function";
  let constructor = "TFDevice::CreateResourceOpLiftingForMainFunctionPass()";
}

def AnnotateParameterReplicationPass :
    Pass<"tf-annotate-parameter-replication", "ModuleOp"> {

    }

    public <T> T realizeModelElement(String path, ModelType<T> type) {
        project.modelRegistry.realize(path, type)
    }

    ModelMap<BinarySpec> getBinaries() {
        realizeModelElement("binaries", modelMap(BinarySpec))
    }

    NativeToolChainRegistry getToolChains() {
        realizeModelElement("toolChains", NativeToolChainRegistry)
    }

  registry.insert<mlir::shape::ShapeDialect>();
  mlir::MLIRContext context(registry);
  mlir::OwningOpRef<mlir::ModuleOp> module;
  {
    mlir::StatusScopedDiagnosticHandler diagnostic_handler(&context);
    module = mlir::parseSourceString<mlir::ModuleOp>(mlir_txt, &context);
    if (!module) {
      tsl::Set_TF_Status_from_Status(status,

        e.cause.message == """Type $ModelMap.name<?> is not a valid model element type:
- type parameter of org.gradle.model.ModelMap cannot be a wildcard.

The type was analyzed due to the following dependencies:
${fullyQualifiedNameOf(managedType)}
  \\--- property 'managedWithNestedInvalidManagedType' (${fullyQualifiedNameOf(nestedManagedType)})
    \\--- property 'invalidManaged' ($ModelMap.name<?>)"""

        where:

constexpr StringRef kCpuDeviceName = "CPU";

using ::mlir::TFL::common::OpsAdded;
using ::mlir::TFL::common::Subgraph;

class RaiseTargetSubgraphsPass
    : public PassWrapper<RaiseTargetSubgraphsPass, OperationPass<ModuleOp>> {
 public:
  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(RaiseTargetSubgraphsPass)

  RaiseTargetSubgraphsPass() = default;
  RaiseTargetSubgraphsPass(const RaiseTargetSubgraphsPass& other) {

};

LogicalResult DecomposeStackOpsInternal(
    Block*, ModuleOp, llvm::SmallDenseMap<Value, Value>*,
    llvm::StringMap<PartitionedCallStackOpsInfo>*);

// Handles stack usage by a tf.While. It will convert the body and conditional
// function signatures, and performs stack ops decomposition on them.
LogicalResult HandleWhileOp(
    TF::WhileOp while_op, ModuleOp module,