From dcf4682be061ecac26a2cdbb4ff96f716772e154 Mon Sep 17 00:00:00 2001
From: superjom <superjom@gmail.com>
Date: Fri, 22 Sep 2017 20:15:49 -0400
Subject: [PATCH] add LoDTensor::NumElements(id,id)

---
 paddle/framework/lod_tensor.cc      |  6 ++++++
 paddle/framework/lod_tensor.h       | 12 ++++++++++++
 paddle/framework/lod_tensor_test.cc |  6 ++++++
 3 files changed, 24 insertions(+)

diff --git a/paddle/framework/lod_tensor.cc b/paddle/framework/lod_tensor.cc
index 88b7f9404a3..513e63657b1 100644
--- a/paddle/framework/lod_tensor.cc
+++ b/paddle/framework/lod_tensor.cc
@@ -77,9 +77,15 @@ size_t LoDTensor::NumElements(size_t level, size_t idx) const {
   PADDLE_ENFORCE_LT(idx, NumElements(level));
   // the last level of LoD, just return number of records in Tensor
   if (level == NumLevels() - 1) {
+    return lod_[level][idx + 1] - lod_[level][idx];
   }
   // high level of LoD, and there is another lower level, return number of
   // lower-level elements
+  auto tmp = SliceInLevel(lod_, level, idx, idx + 1);
+  PADDLE_ENFORCE_GE(tmp.size(), 2);
+  // there is a 0 as a placeholder stored in LoD, so the number of elements
+  // equals lod.size() - 1
+  return tmp[1].size() - 1;
 }
 
 void LoDTensor::SliceLevels(size_t level_begin, size_t level_end) {
diff --git a/paddle/framework/lod_tensor.h b/paddle/framework/lod_tensor.h
index 52f29fb5a38..176e1de4d4b 100644
--- a/paddle/framework/lod_tensor.h
+++ b/paddle/framework/lod_tensor.h
@@ -38,6 +38,18 @@ using Vector = thrust::host_vector<
     T, thrust::system::cuda::experimental::pinned_allocator<T>>;
 #endif
 
+/*
+ * 3-level LoD stores
+ *
+ * 0 10 20
+ * 0 5 10 15 20
+ * 0 2 5 7 10 12 15 20
+ *
+ * - in a level, each element indicates offset in the underlying Tensor
+ * - the first element should be 0 and that indicates that this sequence start
+ * from 0
+ * - each sequence's begin and end(no-inclusive) is level[id, id+1]
+ */
 using LoD = std::vector<Vector<size_t>>;
 
 LoD SliceLevels(const LoD& in, size_t level_begin, size_t level_end);
diff --git a/paddle/framework/lod_tensor_test.cc b/paddle/framework/lod_tensor_test.cc
index 7915326b27a..86db9533cfc 100644
--- a/paddle/framework/lod_tensor_test.cc
+++ b/paddle/framework/lod_tensor_test.cc
@@ -56,6 +56,12 @@ TEST_F(LoDTensorTester, NumElements) {
   ASSERT_EQ(lod_tensor_.NumElements(2), 8UL);
 }
 
+TEST_F(LoDTensorTester, NumElements2) {
+  ASSERT_EQ(lod_tensor_.NumElements(0, 0), 2UL);
+  ASSERT_EQ(lod_tensor_.NumElements(0, 1), 2UL);
+  ASSERT_EQ(lod_tensor_.NumElements(1, 1), 2UL);
+}
+
 TEST_F(LoDTensorTester, SliceLevels) {
   // slice 1 level
   for (size_t level = 0; level < 3UL; ++level) {
-- 
GitLab