#include "tensorflow/core/framework/function.pb.h"
#include "tensorflow/core/protobuf/saved_object_graph.pb.h"

namespace tensorflow {

// This is the TF eager runtime implementation of SignatureDefFunction (separate
// from the TFRT implementation). The user-facing API of SignatureDefFunctions
// and their semantic differences from ConcreteFunction are described here:

  ~AbstractTensorHandle() override {}

 public:
  // Returns tensor dtype.
  virtual tensorflow::DataType DataType() const = 0;

  // Returns the status of the tensor handle. If it is a tfrt::TensorHandle,
  // the tensor handle can be an error and return non-OK status.
  virtual tensorflow::Status TensorHandleStatus() const;

  // Returns tensor shape. If tensor has unknown rank, shape remains untouched.

#ifdef TF_MLRT_OPS
#else
#define TF_MLRT_OPS

include "mlir/Interfaces/SideEffectInterfaces.td"
include "tensorflow/compiler/mlir/tfrt/ir/mlrt/tf_mlrt_dialect.td"
include "tensorflow/compiler/mlir/tfrt/ir/mlrt/mlrt_dialect.td"

def CreateOp: TensorflowMlrt_Op<"createop", []> {
  let summary = "The Fallback CreateOp";

  let description = [{

    const FunctionLibraryRuntime& function_library, DeviceBase* device,
    se::Stream* stream, const XlaPlatformInfo& platform_info,
    bool has_ref_vars);

// Returns created options for XLA compiler when TFRT-TPU is used.
XlaCompiler::Options GenerateCompilerOptionsForTfrtTpu(
    const DeviceCompiler<xla::LocalExecutable, xla::LocalClient>&
        xla_device_compiler,
    const FunctionLibraryRuntime& function_library);

#include "tensorflow/core/framework/device.h"
#include "tensorflow/core/framework/tensor.h"
#include "tensorflow/core/framework/tensor_shape.h"
#include "tensorflow/core/framework/types.h"
#include "tensorflow/core/tfrt/common/pjrt_util.h"
#include "tsl/platform/statusor.h"

namespace tensorflow {
namespace {

TEST(PjRtTensorBufferUtilTest, MakeTensorFromPjRtBuffer) {
  DeviceSetup device_setup;

#include "tensorflow/compiler/mlir/tensorflow/translate/export_tf_dialect_op.h"
#include "tensorflow/compiler/mlir/tensorflow/utils/convert_tensor.h"
#include "tensorflow/core/tfrt/fallback/fallback_state.h"
#include "tensorflow/core/tfrt/fallback/op_kernel_runner.h"

namespace mlir {
namespace TF {

using tensorflow::tfrt_stub::FallbackState;
using tensorflow::tfrt_stub::OpKernelRunner;

#include "tensorflow/core/framework/device.h"
#include "tensorflow/core/framework/device_factory.h"
#include "tensorflow/core/profiler/lib/traceme.h"
#include "tensorflow/core/tfrt/common/async_value_tensor.h"
#include "tensorflow/core/tfrt/common/create_pjrt_client_util.h"
#include "tsl/framework/device_id_utils.h"

namespace tensorflow {
namespace {

absl::StatusOr<std::unique_ptr<xla::PjRtBuffer>> HostTensorToPjRtBuffer(

#include "xla/shape.h"
#include "xla/stream_executor/device_memory.h"
#include "xla/stream_executor/stream.h"
#include "tensorflow/core/framework/device_base.h"
#include "tfrt/concurrency/async_value_ref.h"  // from @tf_runtime

namespace tensorflow {

// XlaHostSendDeviceContext is a DeviceContext that is intended to be
// used to transfer from host->device using Rendezvous. It transfers the

// RUN: tf-opt %s -tf-replicated-clustering-bridge-v2 -tfrt-lower-cluster-to-runtime-ops-tpu 2>&1 | FileCheck %s

// TPUReshardVariables should be inserted even when While functions' shapes are
// different than While operand shapes. Test the whole tf-tpu-bridge because
// correct insertion of TPUReshardVariables depends on multiple passes including
// TPUVariableRuntimeReformatting, ShapeInference, WhileRegion canonicalization,
// and TPUMergeVariablesWithExecute.

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

// CHECK-LABEL: @main
// CHECK-SAME: ([[f0:%.*]]: !mlrt.future, [[f1:%.*]]: !mlrt.future, [[f2:%.*]]: !mlrt.future)
func.func @main(%f0: !mlrt.future, %f1: !mlrt.future, %f2: !mlrt.future) -> (!tf_mlrt.tensor, !tf_mlrt.tensor, !tf_mlrt.tensor) {
  // CHECK-NEXT: [[t:%.*]]:3 = tf_mlrt.await_all [[f0]], [[f1]], [[f2]]
  // CHECK-NOT: tf_mlrt.await
  // CHECK-NEXT: return [[t]]#0, [[t]]#1, [[t]]#2