#include "tensorflow/core/framework/tensor.h"
#include "tensorflow/core/framework/types.h"
#include "tensorflow/core/platform/refcount.h"
#include "tensorflow/core/platform/status.h"
#include "tensorflow/core/tfrt/common/pjrt_util.h"

TF_Device* TF_GetDevice(TF_OpKernelContext* ctx) {
  auto* cc_ctx = reinterpret_cast<tensorflow::OpKernelContext*>(ctx);
  return reinterpret_cast<TF_Device*>(cc_ctx->device());
}

#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 {

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

#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"

TFE_Context* TFE_NewContext(const TFE_ContextOptions* opts, TF_Status* status) {
  if (opts->use_tfrt) {
    status->status = tensorflow::errors::Unimplemented("TFRT is not supported");
    return nullptr;
  }
  std::vector<std::unique_ptr<tensorflow::Device>> devices;
  status->status = tensorflow::DeviceFactory::AddDevices(

  pm.addPass(mlir::TF::CreateTFShapeInferencePass());

  if (lower_to_xla_hlo) {
    // Legalize any StableHLO ops to MHLO. Bridge still doesn't use StableHLO
    // but such ops might be present in the input from upstream like TFRT
    // compilation. Later on, this could be merged in the legalization pass when
    // we migrate bridge to StableHLO.

    // TODO(b/259459405): Avoid this peculiar use through some refactoring in
    // the caller.