absl::StrJoin(missing, ","));
  }
  return absl::OkStatus();
}

void DeviceSetup::AddDevicesAndSetUp(
    const std::vector<std::string>& device_names,
    const std::optional<FunctionDef>& fdef) {
  SessionOptions options;
  auto* device_count = options.config.mutable_device_count();
  for (const auto& device_name : device_names) {
    device_count->insert({device_name, 1});
  }

//
// SavedModelAPI's ConcreteFunctions' lifetimes are bound to the SavedModel they
// are loaded from, since they retain pointers to the TensorHandles owned by the
// SavedModel, and the FunctionDef of the SavedModel.
//
// Note(bmzhao): This class is only TEMPORARILY virtual, as a way to unblock
// TFRT integration with TF Serving. Do not add more virtual implementations of

  // available to be called/referenced by its name in this context.
  virtual Status RegisterFunction(AbstractFunction*) = 0;
  // Remove a function. 'func' argument is the name of a previously added
  // FunctionDef. The name is in fdef.signature.name.
  virtual Status RemoveFunction(const string& func) = 0;

 private:
  const AbstractContextKind kind_;
};

namespace internal {
struct AbstractContextDeleter {

    core::RefCountPtr<AbstractFunction> scoped_func;
    // Returning null tensors from a tf.function is not supported, so we keep
    // track of indices in the model's outputs are nullptr in this set.
    // The FunctionDef only outputs the non-null tensors. We later pad the
    // function op outputs to have nullptrs at the `null_indices`.
    absl::flat_hash_set<int> null_indices;
    {
      AbstractContextPtr func_ctx(BuildFunction(fn_name));

                           const std::string &pass_pipeline,
                           bool show_debug_info, TF_Status *status);

// Simple wrapper to support tf.mlir.experimental.convert_function.
// Load FunctionDef (binary or textual proto format), convert to MLIR, and
// (optionally) optimize the module before returning it as a string.
// This is an early experimental API, ideally we should return a wrapper object

// This pass looks at the graph and all associated FunctionDefs, and turns
// traditional control flow structure (Switch/Merge/etc.) into functional
// control flow structure (XlaIf/XlaWhile). Following passes must
// handle those FunctionDef correctly.

// POST_REWRITE_FOR_EXEC passes that support auto-clustering to enable XLA:

REGISTER_OPTIMIZATION(OptimizationPassRegistry::POST_REWRITE_FOR_EXEC, 5,

}

TEST(ForceXlaConstantsOnHostPassTest, Simple) {
  GraphDefBuilder b(GraphDefBuilder::kFailImmediately);
  Scope root = Scope::NewRootScope().ExitOnError();
  FunctionDefLibrary library;

  FunctionDef called_func =
      FunctionDefHelper::Create("TransposeCall",
                                /*in_def=*/{"a:float", "b:int32"},
                                /*out_def=*/{"c:float"}, {},

    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();
  Fprint128 cache_key =
      orig_op->MutableAttrs()->CacheKey(orig_op->DeviceName());

    return tensorflow::Status(absl::StatusCode::kInternal,
                              absl::StrCat("Failed to find node named ", name));
  }
  return iter->second;
}

const RepeatedPtrField<FunctionDef>& TFPackage::GetFunctionDefs() {
  auto& function_library =
      saved_model_proto_.mutable_meta_graphs(0)->graph_def().library();
  return function_library.function();
}

}  // namespace libexport

    node.mutable_attr()->insert({attr.first, attr.second.proto});
  return node;
}

}  // namespace

GraphDef DeviceCompilerSerializeTest::GetTestGraph(
    const PartialTensorShape& input_shape) {
  FunctionDef make_test_fn = FunctionDefHelper::Define(
      "TestFn", {"a:float", "b:float", "c:float"}, {"m:float"}, {},
      {{{"d"}, "Add", {"a", "b"}, {{"T", DT_FLOAT}}},
       {{"e"}, "Mul", {"d", "c"}, {{"T", DT_FLOAT}}},