add begin/end to ExecuTorch pytree::arr

Differential Revision: D68166302

Pull Request resolved: #7653

Author: swolchok

extension/pytree/aten_util/ivalue_util.cpp

@@ -131,8 +131,8 @@ std::pair<std::vector<at::Tensor>, std::unique_ptr<TreeSpec<Empty>>> flatten(
   auto p = flatten(c);
 
   std::vector<at::Tensor> tensors;
-  for (int i = 0; i < p.first.size(); ++i) {
-    tensors.emplace_back(p.first[i]->toTensor());
+  for (const auto& item : p.first) {
+    tensors.emplace_back(item->toTensor());
   }
 
   return {tensors, std::move(p.second)};

extension/pytree/pytree.h

@@ -431,6 +431,22 @@ struct arr {
     return data_.get();
   }
 
+  T* begin() {
+    return data_.get();
+  }
+
+  T* end() {
+    return begin() + size();
+  }
+
+  const T* begin() const {
+    return data_.get();
+  }
+
+  const T* end() const {
+    return begin() + size();
+  }
+
   inline size_t size() const {
     return n_;
   }

extension/pytree/test/test_pytree.cpp

@@ -11,13 +11,27 @@
 #include <gtest/gtest.h>
 #include <string>
 
+using ::executorch::extension::pytree::arr;
 using ::executorch::extension::pytree::ContainerHandle;
 using ::executorch::extension::pytree::Key;
 using ::executorch::extension::pytree::Kind;
 using ::executorch::extension::pytree::unflatten;
 
 using Leaf = int32_t;
 
+TEST(PyTreeTest, ArrBasic) {
+  arr<int> x(5);
+  ASSERT_EQ(x.size(), 5);
+  for (int ii = 0; ii < x.size(); ++ii) {
+    x[ii] = 2 * ii;
+  }
+  int idx = 0;
+  for (const auto item : x) {
+    EXPECT_EQ(item, 2 * idx);
+    ++idx;
+  }
+}
+
 TEST(PyTreeTest, List) {
   Leaf items[2] = {11, 12};
   std::string spec = "L2#1#1($,$)";