diff options
| author |  Akshay Modi <nareshmodi@google.com> | 2018-09-17 18:42:45 -0700 |
|---|---|---|
| committer |  TensorFlower Gardener <gardener@tensorflow.org> | 2018-09-17 18:46:51 -0700 |
| commit | 3b7ca4b86416f6b6153de90bc1df6e6e5b41934c (patch) | |
| tree | e1eb70719e1b8b4edc572d02e18bcedbd8c00d9b /tensorflow/c | |
| parent | 71fab28dc4741dedf13fea732f6b134608719bc7 (diff) | |
Num elements fastpath for eager tensors.
PiperOrigin-RevId: 213377426
Diffstat (limited to 'tensorflow/c')
| -rwxr-xr-x | tensorflow/c/eager/c_api.cc | 11 | ||||
| -rwxr-xr-x | tensorflow/c/eager/c_api.h | 2 |
2 files changed, 13 insertions, 0 deletions
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, |