#include "tensorflow/core/framework/types.h"
#include "tensorflow/core/lib/core/status.h"
#include "tensorflow/core/platform/mutex.h"
#include "tensorflow/core/platform/stream_executor_no_cuda.h"
#include "tensorflow/core/tfrt/common/async_value_tensor.h"

namespace tensorflow {

class XlaDevice : public LocalDevice {
 public:
  // Given a tensor, sets `xla::Shape*` the shape of tensor's representation

// Cluster groups together operations in the single basic block based on the
// given clustering policy set. Clusters can be outlined into nested modules
// later device specific compilation (e.g. for TFRT JIT compiler).
struct Cluster {
  llvm::SmallVector<Operation*> operations;
  ValuesConstraintSet constraints;
};

// Returns clusters of operations in the given `block` based on the provided

    return stage == IrExportStage::HLO ||
           stage == IrExportStage::HLO_NO_METADATA ||
           stage == IrExportStage::HLO_SERIALIZED;
  };
  // TODO(b/238830423): support GetCompilerIr on TFRT TPU device for stages
  // that requires compilation from HLO to executable.
  if (device_type != DEVICE_CPU && stream == nullptr &&
      !is_tfrt_tpu_supported_stage(stage)) {
    return absl::InternalError(

// RUN: tf-tfrt-opt -split-input-file -tf-mlrt-parallelization %s | FileCheck %s --dump-input=fail --dump-input-filter=all

// CHECK-LABEL: func private @main_stream_{{[0-9]*}}
// CHECK-SAME: ({{%.*}}: tensor<i32>, [[PROMISE:%.*]]: !mlrt.promise)
// CHECK: tf.Sub
// CHECK: tf.Sub
// CHECK: tf.Sub
// CHECK: [[RES:%.*]] = "tf.Sub"
// CHECK: "tf_mlrt.tf_promise"([[PROMISE]], [[RES]])
// CHECK: return

// CHECK-LABEL: func @main

#include "tensorflow/core/platform/stream_executor_no_cuda.h"
#include "tensorflow/core/profiler/lib/traceme.h"
#include "tensorflow/core/public/session_options.h"
#include "tensorflow/core/public/version.h"
#include "tensorflow/core/tfrt/common/async_value_tensor.h"
#include "tensorflow/core/util/device_name_utils.h"
#include "tensorflow/core/util/dump_graph.h"
#include "tensorflow/core/util/stream_executor_util.h"

namespace tensorflow {

  if (TF_GetCode(status) != TF_OK) return;

  // unwrap op_ and set step_id only if valid step id value was set.
  // Currently only required for non-TFRT use cases, e.g., EagerOp.
  if (step_id_.has_value()) {
    tensorflow::unwrap(op_.get())->SetStepId(step_id_.value());
  }

  TFE_Execute(op_.get(), unwrapped_results.data(), &real_num_outputs, status);

// RUN: tf-tfrt-opt -split-input-file -tf-to-mlrt %s | FileCheck %s

// CHECK-LABEL: @main_stream_0
// CHECK-SAME: ([[input0:%.*]]: !tf_mlrt.tensor, [[promise_b:%.*]]: !mlrt.promise)
func.func @main_stream_0(%input0: tensor<i32>, %promise_b: !mlrt.promise) {
  %const = "tf.Const"() {__op_key = 0 : i32, value = dense<1> : tensor<i32>} : () -> tensor<i32>
  // CHECK: [[a:%.*]] = tf_mlrt.executeop([[input0]],
  // CHECK-SAME: AddV2

// RUN: tf-tfrt-opt -split-input-file -tf-mlrt-async-while %s | FileCheck %s

// This is a simple case that should be pipelined.

// CHECK-LABEL: func.func private @"map/while_cond"
func.func private @"map/while_cond"(%loop_count: tensor<i32>, %max_iterations: tensor<i32>, %handle: tensor<?x!tf_type.resource>, %flow_in: tensor<*xf32>, %matrix: tensor<3x3xf32>) -> tensor<i1> {
  %0 = "tf.Less"(%loop_count, %max_iterations) : (tensor<i32>, tensor<i32>) -> tensor<i1>

/// derived from the `Handle` class. Instances of `Handle` are referred to as
/// "handles". All handles have a tagged value.
///
/// Example Usage:
/// Object runtime = GetRuntime("tfrt");
/// Object module = runtime.Get("Import")("cool_mobilenet")
/// runtime.Get("Tensor")(Tuple(5,5,5), 3.3);
/// Object test = CreateModule("test");
/// test.Set("cool_function", callable);

#include "tensorflow/core/lib/core/refcount.h"
#include "tensorflow/core/platform/errors.h"
#include "tensorflow/core/platform/status.h"
#include "tensorflow/core/platform/statusor.h"
#include "tensorflow/core/tfrt/common/async_value_tensor.h"
#include "tensorflow/core/util/stream_executor_util.h"
#include "tsl/framework/device_id_utils.h"
#include "tsl/framework/serving_device_selector_policies.h"
#include "tsl/platform/status.h"