return registry;
        }()) {
    context_.loadAllAvailableDialects();

    mlir::OpBuilder builder(&context_);
    module_ = builder.create<mlir::ModuleOp>(builder.getUnknownLoc());

    builder.setInsertionPointToStart(module_->getBody());
    auto func = builder.create<mlir::func::FuncOp>(  //
        builder.getUnknownLoc(), "main", builder.getFunctionType({}, {}));
    func->setAttr("tfl.func", builder.getUnitAttr());

namespace {

// Helper class for parsing operands to a foldable operation.
class FoldAdaptor {
 public:
  // Returns std::nullopt if the operation cannot be folded.
  static std::optional<FoldAdaptor> Create(Operation* operation) {
    auto foldable_opr = [](Value val) -> bool {
      return !llvm::isa<BlockArgument>(val) &&
             llvm::isa<stablehlo::ConstantOp>(val.getDefiningOp());
    };

  TensorShape shape(shape_vector);

  // Create the variable.
  Status status;
  std::unique_ptr<Variable> var;
  TF_EXPECT_OK(Variable::CreateUninitialized(context(), dtype, shape,
                                             absl::nullopt, nullptr, {}, &var));

  // Create a TensorHandle
  ImmediateTensorHandlePtr expected_handle =

    bool used_persistent_cache) {
  metrics::UpdateXlaCompilationTime(compile_time_us);

  const std::string& function_name = function.name();

  mutex_lock lock(mu_);
  // Create a stats entry if it doesn't already exist.
  auto it =
      cluster_compile_stats_.emplace(function.name(), ClusterCompileStats{})
          .first;

  const uint64 compile_time_s = compile_time_us / 1.0e6;

  core::RefCountPtr<Var> variable;
  const Tensor& value = context->input(1);
  // Note: every resource-variable-manipulating op assumes copy-on-write
  // semantics, and creates a copy of the variable's Tensor if its refcount is
  // bigger than 1 when we try to modify it. This means we never need to copy
  // the original tensor for AssignVariableOp; even if there are other live

                                    std::move(patterns)))) {
    getOperation().emitError("Composite lowering pass failed.");
    signalPassFailure();
  }
}

}  // namespace

// Creates an instance of the pass.
std::unique_ptr<OperationPass<ModuleOp>> CreateCompositeLoweringPass() {
  return std::make_unique<CompositeLoweringPass>();
}

// Registers the pass implementation

  } else if (op_.NumOutputs() == 1) {
    RenderExecutionSingleOutput();
  } else {
    RenderExecutionMultipleOutputs();
  }
}

void OpImplementationRenderer::RenderInitialization() {
  // Create Op variable and initialize it
  Statement("AbstractOperationPtr $0(ctx->CreateOperation())",
            op_.VariableName());
  TFStatement(Call(op_.VariableName(), "Reset",

  // context in tensorflow_accelerator_device_info(). Also,
  // tensorflow_accelerator_device_info() == nullptr is used as an IsCPU test.
  // We need XlaCpuDevice to be treated not as CPU because it allocates
  // XlaTensors, not regular Tensors.
  Status status = device->UseAcceleratorDeviceInfo();
  if (!status.ok()) {
    errors::AppendToMessage(&status, "while setting up ", DEVICE_GPU_XLA_JIT);

      return op->emitError("expected exactly one called_computation");

    auto decomposition = mlir::cast<FlatSymbolRefAttr>(calledComputations[0]);

    auto composite = rewriter.create<mlir::stablehlo::CompositeOp>(
        op.getLoc(), op.getResultTypes(), op.getOperands(), name.str(), attrs,
        decomposition.getValue());
    rewriter.replaceOp(op, composite.getResults());
    return success();
  }

    // the resource, and we use it to just look up.
    // See core/kernels/data/iterator_ops.cc.)
    resource_op = builder.create<TF::IteratorOp>(
        op->getLoc(), resource.getType(), shared_name, container,
        it.getOutputTypes(), it.getOutputShapes());
  } else {
    resource_op = builder.create<TF::VarHandleOp>(
        op->getLoc(), resource.getType(), container, shared_name);
  }