TFE_Context* ctx) {
  return static_cast<TFE_ContextDevicePlacementPolicy>(
      tensorflow::unwrap(ctx)->GetDevicePlacementPolicy());
}

TFE_TensorHandle* TFE_NewTensorHandle(const TF_Tensor* t, TF_Status* status) {
  tensorflow::Tensor tensor;
  status->status = tensorflow::TF_TensorToTensor(t, &tensor);
  if (!status->status.ok()) return nullptr;

  }

  auto* retval_task0 = TFE_TensorHandleCopyToDevice(
      retvals[0], ctx, "/job:localhost/replica:0/task:0/device:CPU:0", status);
  ASSERT_EQ(TF_GetCode(status), TF_OK) << TF_Message(status);

  TF_Tensor* t = TFE_TensorHandleResolve(retval_task0, status);
  ASSERT_EQ(TF_GetCode(status), TF_OK) << TF_Message(status);
  TFE_DeleteTensorHandle(retval_task0);
  float product[4] = {0};

template <typename value_type>
void ExpectScalarEq(TFE_TensorHandle* handle, value_type expected_value) {
  std::unique_ptr<TF_Status, decltype(&TF_DeleteStatus)> status(
      TF_NewStatus(), TF_DeleteStatus);
  std::unique_ptr<TF_Tensor, decltype(&TF_DeleteTensor)> actual_value(
      TFE_TensorHandleResolve(handle, status.get()), TF_DeleteTensor);
  ASSERT_EQ(TF_GetCode(status.get()), TF_OK) << TF_Message(status.get());

template <class T, TF_DataType datatype>
TFE_TensorHandle* TestTensorHandleWithDims(TFE_Context* ctx, const T* data,
                                           const int64_t* dims, int num_dims) {
  TF_Status* status = TF_NewStatus();
  TF_Tensor* t = TFE_AllocateHostTensor(ctx, datatype, dims, num_dims, status);
  memcpy(TF_TensorData(t), data, TF_TensorByteSize(t));
  TFE_TensorHandle* th = TFE_NewTensorHandleFromTensor(ctx, t, status);

  if (t == nullptr) {
    status->status =
        tensorflow::errors::InvalidArgument("Unsupported dtype: ", dtype);
    return nullptr;
  }

  return new TF_Tensor{t};
}

TFE_TensorHandle* TFE_NewTensorHandleFromTensor(TFE_Context* ctx, TF_Tensor* t,
                                                TF_Status* status) {
  return tensorflow::wrap(
      tensorflow::unwrap(ctx)->CreateLocalHandle(t->tensor));
}

  // appended later), but should include any information specific to this custom
  // device which would be useful for debugging.
  //
  // Optional. If null, defaults to resolving the TFE_TensorHandle into a
  // TF_Tensor and summarizing that.
  TF_Buffer* (*summarize)(void* data, TF_Status* status) = nullptr;
} TFE_CustomDeviceTensorHandle;

// Creates a new TensorHandle from memory residing in a custom device. Takes

        /*attr=*/{{"T", tensorflow::DataTypeSlice({DT_FLOAT, DT_FLOAT})}},
        /*dep=*/{},
        /*device=*/""}});
}

template <typename T>
void ExpectEquals(AbstractTensorHandle* t, T expected) {
  TF_Tensor* result_t;
  Status s = tensorflow::GetValue(t, &result_t);
  ASSERT_TRUE(s.ok()) << s.message();
  auto value = static_cast<T*>(TF_TensorData(result_t));
  EXPECT_EQ(*value, expected);
  TF_DeleteTensor(result_t);

// CHECK-LABEL: func @main
// CHECK-SAME: {{.*}} !tfrt.chain
// CHECK-SAME: [[serialized:%.*]]: !tfrt_fallback.tf_tensor
func.func @main(%serialized: tensor<32x!tf_type.string>) -> (tensor<?x2xi64>) attributes {tf.entry_function = {inputs = "input0", outputs = "ParseExample/ParseExampleV2"}} {

  TFE_TensorHandle* var_value_unpacked =
      UnpackTensorHandle(var_value, status.get());
  ASSERT_TRUE(TF_GetCode(status.get()) == TF_OK) << TF_Message(status.get());
  std::unique_ptr<TF_Tensor, decltype(&TF_DeleteTensor)> resolved_value(
      TFE_TensorHandleResolve(var_value_unpacked, status.get()),
      TF_DeleteTensor);
  ASSERT_TRUE(TF_GetCode(status.get()) == TF_OK) << TF_Message(status.get());

  for (int device_index = 0; device_index < num_underlying_devices();
       ++device_index) {
    auto device_value = absl::make_unique<DataType>();
    *device_value = values[device_index];
    std::unique_ptr<TF_Tensor, decltype(&TF_DeleteTensor)> tensor(
        TF_NewTensor(
            datatype_enum, /*dims=*/nullptr, /*num_dims=*/0,
            device_value.release(), sizeof(DataType),
            [](void* data, size_t, void* arg) {