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");

   `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.

 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

    //   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,

                        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);

namespace tensorflow {

namespace {
using DeadnessPredicate = DeadnessAnalysis::DeadnessPredicate;
using jit::DeviceId;
using jit::DeviceSet;

// The clusters we create here are eventually lowered into an
// _XlaCompile/_XlaRun pair with a TF executor "fallback" that uses the
// PartitionedCall op to execute the cluster in the regular graph executor if
// need be.  PartitionedCall, however, reruns the entire TF graph optimization

        &effects) {
  effects.reserve(2 * getArgs().size() + 1);
  effects.emplace_back(MemoryEffects::Write::get(),
                       ResourceEffects::_XlaRun::get());

  // Conservatively mark resource handles as read and write, as without
  // analyzing _XlaCompile, there is not sufficient information to determine
  // effects on resources.

  let summary = "XLA Compile Op. For use by the XLA JIT only.";

  let description = [{
Compiles a TensorFlow function into an XLA LocalExecutable and returns a key
that _XlaRun can use to look up the LocalExecutable and execute it.
  }];

  let arguments = (ins
    Variadic<TF_Tensor>:$constants,
    Variadic<TF_Tensor>:$args,
    Variadic<TF_ResourceTensor>:$resources,