From 3b7ca4b86416f6b6153de90bc1df6e6e5b41934c Mon Sep 17 00:00:00 2001 From: Akshay Modi Date: Mon, 17 Sep 2018 18:42:45 -0700 Subject: Num elements fastpath for eager tensors. PiperOrigin-RevId: 213377426 --- tensorflow/c/eager/c_api.cc | 11 +++++++++++ tensorflow/c/eager/c_api.h | 2 ++ 2 files changed, 13 insertions(+) (limited to 'tensorflow/c') diff --git a/tensorflow/c/eager/c_api.cc b/tensorflow/c/eager/c_api.cc index 6f86ea80e5..0bf3d9542b 100755 --- a/tensorflow/c/eager/c_api.cc +++ b/tensorflow/c/eager/c_api.cc @@ -375,6 +375,17 @@ int TFE_TensorHandleNumDims(TFE_TensorHandle* h, TF_Status* status) { return result; } +int64_t TFE_TensorHandleNumElements(TFE_TensorHandle* h, TF_Status* status) { + if (h == nullptr || h->handle == nullptr) { + status->status = tensorflow::errors::InvalidArgument( + "The passed in handle is a nullptr"); + return -1; + } + tensorflow::int64 result; + status->status = h->handle->NumElements(&result); + return result; +} + int64_t TFE_TensorHandleDim(TFE_TensorHandle* h, int dim_index, TF_Status* status) { if (h == nullptr || h->handle == nullptr) { diff --git a/tensorflow/c/eager/c_api.h b/tensorflow/c/eager/c_api.h index a87d73ec8e..6323f8a053 100755 --- a/tensorflow/c/eager/c_api.h +++ b/tensorflow/c/eager/c_api.h @@ -163,6 +163,8 @@ TF_CAPI_EXPORT extern TF_DataType TFE_TensorHandleDataType(TFE_TensorHandle* h); // This function will block till the operation that produces `h` has completed. TF_CAPI_EXPORT extern int TFE_TensorHandleNumDims(TFE_TensorHandle* h, TF_Status* status); +TF_CAPI_EXPORT extern int64_t TFE_TensorHandleNumElements(TFE_TensorHandle* h, + TF_Status* status); // This function will block till the operation that produces `h` has completed. TF_CAPI_EXPORT extern int64_t TFE_TensorHandleDim(TFE_TensorHandle* h, int dim_index, -- cgit v1.2.3