// "Strict" compilation:  every _XlaCompile invocation must compile the
    // cluster.
    ops::_XlaRun xla_run(root.WithOpName("xla_run"), xla_run_args,
                         xla_compile.key, n->output_types());

    MoveOutgoingEdges(g, /*old_node=*/n,
                      /*new_node=*/xla_run.operation.node());
    g->RemoveNode(n);
  } else {

  auto xla_run =
      NodeWith(Op("_XlaRun"), Inputs(Out(1, predicated_compilation_key)));
  auto tf_call =
      NodeWith(Op("StatefulPartitionedCall"),
               CtrlDeps(NodeWith(Op("Identity"),
                                 Inputs(Out(0, predicated_compilation_key)))));
  auto merge = NodeWith(Op("_XlaMerge"), Inputs(Out(tf_call), Out(xla_run)));

REGISTER_OP("_XlaMerge")
    .Input("partitioned_call: T")
    .Input("xla_run: T")
    .Output("output: T")
    .Attr("T: type")
    .SetShapeFn([](InferenceContext* c) {
      c->set_output(0, c->input(0));
      return absl::OkStatus();
    })
    .Doc(R"(XLA Merge Op. For use by the XLA JIT only.

Merges the outputs from the PartitionedCall node and the _XlaRun node.

    : 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

    }

    // Allow using Device API (PjRt) for `device_type` in the XlaCompile and
    // XlaRun ops. Please note that `enabled_for_compile_and_run_` needs to be
    // true in addition to the `device_type` being allowed in order to use the
    // Device API for single device compilation and execution in the XlaCompile
    // and XlaRun ops.
    void AllowForDeviceInXlaCompileAndRun(const DeviceType& device_type) {

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

Laura Kassovic <******@****.***> 1701107622 -0800

}

}

// -----

module attributes {tf_saved_model.semantics} {

// Test not hoisting callees in xla launch functions.

// CHECK-LABEL: func private @xla_func
func.func private @xla_func(%arg0: tensor<1x3xf32>) -> tensor<1x3xf32>
  attributes {tf._input_shapes = [#tf_type.shape<1x3>, #tf_type.shape<*>], tf.signature.is_stateful} {
  // CHECK-NOT: tf._TfrtGetResource

    These operations do not map 1-1 to TensorFlow ops and requires a lowering
    pass later to transform them into Compile/Run op pairs, like XlaCompile and
    XlaRun.
}];

  let cppNamespace = "::mlir::tf_device";
}

//===----------------------------------------------------------------------===//
// TensorFlow Device Dialect Ops definitions