From f414a7f36a0ca6db966363ea729d827b1c83afee Mon Sep 17 00:00:00 2001
From: Github Executorch <github_executorch@arm.com>
Date: Tue, 14 Jan 2025 11:08:16 -0800
Subject: [PATCH] add begin/end to ExecuTorch pytree::arr

Allows use of range-for.

Differential Revision: [D68166302](https://our.internmc.facebook.com/intern/diff/D68166302/)

[ghstack-poisoned]
---
 extension/pytree/aten_util/ivalue_util.cpp |  4 ++--
 extension/pytree/pytree.h                  | 16 ++++++++++++++++
 2 files changed, 18 insertions(+), 2 deletions(-)

diff --git a/extension/pytree/aten_util/ivalue_util.cpp b/extension/pytree/aten_util/ivalue_util.cpp
index c4d11c13ee0..888c9829523 100644
--- a/extension/pytree/aten_util/ivalue_util.cpp
+++ b/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)};
diff --git a/extension/pytree/pytree.h b/extension/pytree/pytree.h
index 515ba8b7402..b0f2ca1fab6 100644
--- a/extension/pytree/pytree.h
+++ b/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_;
   }