LOG_FIRST_N(INFO, 1) << "Skipping Graph Decomposition Pass, decomposition"
                            " library was not found";
    return absl::OkStatus();
  }

  tf_core_op_expansion_graph_counter->GetCell()->IncrementBy(1);

  LOG_FIRST_N(INFO, 1) << "Run Graph Decomposition Passes";

  TF_RETURN_IF_ERROR(DecomposeGraph(module));

  LOG_FIRST_N(INFO, 1) << "Finish Graph Decomposition Passes";

  };
  for (const char* skip : kOpsToSkip) {
    if (absl::StartsWith(orig_op->op_name(), skip)) return absl::OkStatus();
  }

  tf_core_op_expansion_node_counter->GetCell()->IncrementBy(1);

  LOG_FIRST_N(INFO, 1) << "Run Node Expansion Passes";

  // Get the FunctionDef and insert that into the context
  const NodeDef& ndef = orig_op->MutableAttrs()->BuildNodeDef();
  auto& ctx = orig_op->EagerContext();

  if (overall_state == MlirOptimizationPassState::Disabled) {
    if (VLOG_IS_ON(1)) {
      LOG_FIRST_N(INFO, 1)
          << "None of the MLIR Optimization Passes are enabled "
          << "(registered " << registry_->passes().size() << ")";
    }
    return absl::OkStatus();
  }

  if (VLOG_IS_ON(1)) {
    LOG_FIRST_N(INFO, 1) << "MLIR Graph Optimization Passes."
                         << " Enabled: " << num_passes_enabled

    std::vector<std::vector<xla::Shape>>* per_core_arg_shapes,
    std::vector<std::unique_ptr<mlir::Pass>>& custom_legalization_passes,
    xla::CompileOnlyClient* client, XlaCompilationResult* compilation_result) {
  LOG_FIRST_N(INFO, 1) << "Compiling MLIR computation to XLA HLO using the "
                          "Combined MLIR Tf2Xla Bridge.";

  absl::StatusOr<std::string> mlir_compilation =
      internal::CompileFromMlirToXlaHlo(

    const std::vector<TensorShape>& arg_shapes,
    std::vector<ShardingAndIndex>* arg_core_mapping,
    std::vector<std::vector<xla::Shape>>* per_core_arg_shapes) {
  LOG_FIRST_N(INFO, 1)
      << "Compiling MLIR computation to XLA HLO using MLIR tf2xla bridge in "
         "the op by op fallback mode. This is Phase 2 of the TF2XLA Bridge. "

    if (getenv("MLIR_CRASH_REPRODUCER_DIRECTORY"))
      path = getenv("MLIR_CRASH_REPRODUCER_DIRECTORY");
    else if (getenv("TEST_UNDECLARED_OUTPUTS_DIR"))
      path = "sponge";
  }
  if (path.empty()) {
    LOG_FIRST_N(INFO, 1) << "disabling MLIR crash reproducer, set env var "
                            "`MLIR_CRASH_REPRODUCER_DIRECTORY` to enable.";
    return;
  }

      device_layout);
  if (first_try_buffer.ok()) {
    return std::move(*first_try_buffer);
  }
  if (first_try_buffer.status().code() == absl::StatusCode::kUnimplemented) {
    LOG_FIRST_N(WARNING, 1)
        << first_try_buffer.status()
        << "; fallback to BufferFromHostBuffer without device layout.";
    TF_ASSIGN_OR_RETURN(
        std::unique_ptr<xla::PjRtBuffer> second_try_buffer,

    std::vector<tpu::ShardingAndIndex>* arg_core_mapping,
    std::vector<std::vector<xla::Shape>>* per_core_arg_shapes,
    xla::CompileOnlyClient* client,
    XlaCompiler::CompilationResult* compilation_result) {
  LOG_FIRST_N(INFO, 1) << "Compiling MLIR computation to XLA HLO using the "
                          "old (non-MLIR) tf2xla bridge";

  CompilationTimer timer;
  *compilation_result = {};