See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/

#include "tensorflow/c/eager/abstract_tensor_handle.h"

namespace tensorflow {

std::string AbstractTensorHandle::DebugString() const {
  PartialTensorShape shape;
  Status s = Shape(&shape);
  std::string shape_string;
  if (!s.ok()) {

See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/

#include "tensorflow/c/eager/immediate_execution_tensor_handle.h"

namespace tensorflow {

std::string ImmediateExecutionTensorHandle::DebugString() const {
  PartialTensorShape shape;
  std::string shape_string;
  if (Shape(&shape).ok()) {

absl::StatusOr<PJRT_Buffer*> GetPjRtCBufferFromTensor(const Tensor* tensor) {
  tensorflow::AsyncValueTensor* av_tensor =
      tensorflow::AsyncValueTensor::FromTensor(tensor);
  if (av_tensor == nullptr || av_tensor->GetBuffer() == nullptr) {
    return absl::InternalError("Input tensor does not have PjRtBuffer.");
  }
  auto* c_api_buffer =
      dynamic_cast<xla::PjRtCApiBuffer*>(av_tensor->GetBuffer().get());
  if (c_api_buffer == nullptr) {

  tensorflow::AsyncValueTensor* av_tensor =
      tensorflow::AsyncValueTensor::FromTensor(&tensor);
  av_tensor->SetBuffer(std::move(buffer));

  TF_ASSERT_OK_AND_ASSIGN(auto c_buffer, GetPjRtCBufferFromTensor(&tensor));

  EXPECT_THAT(c_buffer, NotNull());
}

TEST(TensorPjRtBufferUtilTest, SetPjRtCBufferToTensorNotAsyncValueTensor) {
  tensorflow::Tensor tensor(DT_FLOAT, {1});

  const DLTensor* dltensor_out = &dlm_out->dl_tensor;
  EXPECT_EQ(dltensor_out->device.device_type, dltensor_in->device.device_type);
  EXPECT_EQ(dltensor_out->device.device_id, dltensor_in->device.device_id);
  EXPECT_EQ(dltensor_out->ndim, dltensor_in->ndim);
  EXPECT_EQ(dltensor_out->dtype.code, dltensor_in->dtype.code);
  EXPECT_EQ(dltensor_out->dtype.bits, dltensor_in->dtype.bits);

#include "tensorflow/c/experimental/gradients/nn_grad.h"

#include "absl/types/span.h"
#include "tensorflow/c/eager/abstract_tensor_handle.h"
#include "tensorflow/c/eager/immediate_execution_context.h"
#include "tensorflow/c/eager/immediate_execution_tensor_handle.h"
#include "tensorflow/c/experimental/ops/array_ops.h"
#include "tensorflow/c/experimental/ops/math_ops.h"
#include "tensorflow/c/experimental/ops/nn_ops.h"

      numerical_grad.reset(numerical_grad_raw);
    }

    TF_Tensor* numerical_tensor;
    s = GetValue(numerical_grad.get(), &numerical_tensor);
    ASSERT_EQ(errors::OK, s.code()) << s.message();
    auto num_elem_numerical = TF_TensorElementCount(numerical_tensor);

    TF_Tensor* analytical_tensor;
    s = GetValue(outputs[i], &analytical_tensor);
    ASSERT_EQ(errors::OK, s.code()) << s.message();

}
TF_AbstractTensor* TF_OutputListGet(TF_OutputList* o, int i) {
  return wrap(unwrap(o)->outputs[i]);
}
void TF_OutputListPushBack(TF_OutputList* o, TF_AbstractTensor* tensor,
                           TF_Status* s) {
  unwrap(o)->outputs.push_back(unwrap(tensor));
}

void TF_AbstractOpSetOpType(TF_AbstractOp* op, const char* const op_type,
                            TF_Status* s) {
  tsl::Set_TF_Status_from_Status(

  ASSERT_EQ(errors::OK, s.code()) << s.message();

  TF_Tensor* result_tensor;
  s = GetValue(outputs[0], &result_tensor);
  ASSERT_EQ(errors::OK, s.code()) << s.message();
  auto result_value = static_cast<float*>(TF_TensorData(result_tensor));
  EXPECT_EQ(*result_value, 1.0);
  outputs[0]->Unref();
  TF_DeleteTensor(result_tensor);
  result_tensor = nullptr;
}

INSTANTIATE_TEST_SUITE_P(

#include "tensorflow/c/eager/abstract_context.h"
#include "tensorflow/c/eager/abstract_tensor_handle.h"
#include "tensorflow/c/eager/c_api.h"
#include "tensorflow/c/eager/c_api_unified_experimental.h"
#include "tensorflow/c/eager/c_api_unified_experimental_internal.h"
#include "tensorflow/c/eager/immediate_execution_context.h"
#include "tensorflow/c/eager/immediate_execution_tensor_handle.h"
#include "tensorflow/c/eager/tfe_context_internal.h"