// CHECK: "tf.XlaLaunch"(%arg1, %arg0) <{function = @func_with_resources, operandSegmentSizes = array<i32: 0, 1, 1>}> : (tensor<i32>, tensor<!tf_type.resource>) -> tensor<i32>
    %0 = "tf_device.cluster_func"(%arg0, %arg1) {func = @func_with_resources} : (tensor<!tf_type.resource>, tensor<i32>) -> tensor<i32>

  //    functions contain the computations to be passed to XlaLaunch. During
  //    encapsulation, we sort the arguments into the order expected by
  //    XlaLaunch.
  static Status Encapsulate(std::unique_ptr<Graph>* graph,
                            FunctionLibraryDefinition* flib_def);

  // b) we rewrite the function calls generated in phase (a) into XlaLaunch
  //    operators. We also convert the XlaClusterOutput output nodes of the

};

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

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

struct WriteTrainingPredictions
    : public ::mlir::SideEffects::Resource::Base<WriteTrainingPredictions> {

// the arguments into the order expected by XlaLaunch computations:
// 1) arguments
// 2) resource variable arguments
// See the documentation of EncapsulateSubgraphsInFunctions for the meaning
// of the arguments.
//
// TODO(b/113166435): Ordering constraints on XlaLaunch op can be relaxed.
Status RewriteSubgraph(const std::vector<OutputTensor>& arg_source_tensors,

  // 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
  // XlaLaunch/XlaCompileOnDemand op is placed and thus does not die before the
  // op kernel.

  // CHECK-LABEL: @converts_cluster
  func.func @converts_cluster() {
    // CHECK: "tf.XlaLaunch"()

  %1 = "tf.ReadVariableOp"(%0) {device = "/device:CPU:0"} : (tensor<!tf_type.resource<tensor<f32>>>) -> tensor<f32>
  %2 = "tf.XlaLaunch"(%arg0, %1) {_noinline = true, _xla_compile_device_type = "GPU", device = "/device:GPU:0", function = @xla_func_0, operandSegmentSizes = array<i32: 0, 2, 0>} : (tensor<7xf32>, tensor<f32>) -> tensor<10xf32>
  func.return %2 : tensor<10xf32>

    This pass rewrites `tf.PartitionedCall` and `tf.StatefulPartitionedCall`
    operations with `_xla_compile_device_type` attribute in a
    `tf_device.cluster` into `tf.XlaLaunch` operations. This makes the attached
    function execute with XLA. `tf.XlaLaunch` requires resource-type arguments
    come at the end, so this pass rewrites the called function if necessary.
    This pass assumes there are no nested `tf_device.cluster`s so we don't end

#include "absl/status/status.h"
#include "tensorflow/core/framework/op.h"
#include "tensorflow/core/framework/shape_inference.h"

namespace tensorflow {

using shape_inference::InferenceContext;

REGISTER_OP("XlaLaunch")
    .Input("constants: Tconstants")
    .Attr("Tconstants: list(type) >= 0")
    .Input("args: Targs")
    .Attr("Targs: list(type) >= 0")
    .Input("resources: Nresources * resource")

  %1 = "tf.ReadVariableOp"(%0) {device = "/device:CPU:0"} : (tensor<!tf_type.resource<tensor<1x3xf32>>>) -> tensor<1x3xf32>
  %2 = "tf.XlaLaunch"(%arg0, %1) {_noinline = true, _xla_compile_device_type = "GPU", device = "/device:GPU:0", function = @xla_func_0, operandSegmentSizes = array<i32: 0, 2, 0>} : (tensor<1x3xf32>, tensor<1x3xf32>) -> tensor<1x3xf32>