TF_ASSERT_OK(BuildXlaOps(root, fdef_lib, &graph));

  Node* write_op_new = FindNodeByName(graph.get(), write_op->name());
  ASSERT_NE(write_op_new, nullptr);
  EXPECT_THAT(write_op_new, NodeWith(CtrlDeps(NodeWith(Op("_XlaRun")))));
}

TEST_F(BuildXlaOpsTest, CleanFailureOnBogusAttr) {
  Scope root = Scope::NewRootScope().ExitOnError();

  FunctionDefLibrary fdef_lib =
      CreateFunctionDefLibWithConstFunction("cluster_0");

  VariableInfo(const VariableInfo&) = delete;
  VariableInfo& operator=(const VariableInfo&) = delete;

  // The index of the DT_RESOURCE input to the _XlaCompile/_XlaRun operator.
  // Note that the indices can be different between _XlaCompile and _XlaRun.
  int index() const { return index_; }

  // A pointer to the resource variable.  May be null if this VariableInfo is
  // "empty", i.e. it does not track a resource variable.

   `compilation_successful` is always true.
)");

REGISTER_OP("_XlaRun")
    .Input("args: Targs")
    .Attr("Targs: list(type) >= 0")
    .Output("results: Tresults")
    .Attr("Tresults: list(type) >= 0")
    .Input("key: string")
    // XLA random-number generation ops are stateful.
    // TODO(phawkins): create stateful and non-stateful variants of _XlaRun.
    .SetIsStateful()
    .Doc(R"(XLA Run Op. For use by the XLA JIT only.

    : public ::mlir::SideEffects::Resource::Base<WriteTrainingPredictions> {
  StringRef getName() final { return "WriteTrainingPredictions"; }
};

struct _XlaRun : public ::mlir::SideEffects::Resource::Base<_XlaRun> {
  StringRef getName() final { return "_XlaRun"; }
};

// Returns true iff resource type with given ID is only self-dependent, i.e.,
// there are no dependencies to other resource types (including unknown resource

 private:
  DeviceType device_type_;
  se::Platform::Id platform_id_;

  // xla_device_metadata_ lives in the tensorflow::DeviceBase in which the
  // XlaLaunch/_XlaCompile/_XlaRun op is placed and thus does not die before the
  // XlaLaunch/_XlaCompile/_XlaRun OpKernel.
  const XlaDevice::Metadata* xla_device_metadata_;

  // pjrt_device_metadata_ lives in tensorflow::PjRtBaseDevice in which the

#include "tensorflow/core/common_runtime/optimization_registry.h"
#include "tensorflow/core/lib/core/status.h"

namespace tensorflow {

// Replaces TF function calls marked with `_XlaCompiledKernel` with _XlaCompile
// and _XlaRun nodes (which compile and launch, respectively, the corresponding
// HLO module).
class BuildXlaOpsPass : public GraphOptimizationPass {
 public:
  // If enable_lazy_compilation is not nullopt then *enable_lazy_compilation

    //   xla_run_outputs = _XlaRun(..., key=use_xla_run)
    //   outputs = Merge(tf_call_outputs, xla_run_outputs).
    ops::Switch s(root.WithOpName("predicated_compilation_key"),
                  xla_compile.key, xla_compile.compilation_successful);
    Output predicated_compilation_key = s.output_true;
    Output inverse_predicated_compilation_key = s.output_false;

    ops::_XlaRun xla_run(root.WithOpName("xla_run"), xla_run_args,

                              .HostMemory("resources"),             \
                          KERNEL);

#define REGISTER_XLA_RUN_KERNEL(DEVICE, KERNEL, TYPES) \
  REGISTER_KERNEL_BUILDER(Name("_XlaRun").Device(DEVICE), KERNEL);

#define REGISTER_XLA_DEVICE_KERNELS(DEVICE, TYPES)                             \
  REGISTER_KERNEL_BUILDER(                                                     \

                        XlaCompileOp);

REGISTER_KERNEL_BUILDER(Name("_XlaRun").Device(DEVICE_CPU), XlaRunOp);
REGISTER_KERNEL_BUILDER(Name("_XlaRun").Device(DEVICE_GPU).HostMemory("key"),
                        XlaRunOp);
REGISTER_KERNEL_BUILDER(
    Name("_XlaRun").Device(DEVICE_DEFAULT).HostMemory("key"), XlaRunOp);

def TF_NcclAllReduceOrderingResource : TF_ResourceBase<"NcclAllReduceOrdering">;
def TF_GlobalIterIdResource : TF_ResourceBase<"GlobalIterId">;
def TF__XlaRunResource : TF_ResourceBase<"_XlaRun">;
// Fake resource, see `TF_MustExecute` below.
def TF_MustExecuteResource : TF_ResourceBase<"MustExecute">;

// Value-based side effects
//