Obtained from:	https://github.com/metaopt/optree/commit/9a613a705ee18e18784ba3a16b2786f56f9c74d2 (based on)

--- include/treespec.h.orig	2024-03-25 18:01:43 UTC
+++ include/treespec.h
@@ -19,6 +19,7 @@ limitations under the License.
 
 #include <pybind11/pybind11.h>
 
+#include <algorithm>      // std::min
 #include <memory>         // std::unique_ptr
 #include <optional>       // std::optional, std::nullopt
 #include <string>         // std::string
@@ -38,7 +39,10 @@ using ssize_t = py::ssize_t;
 using ssize_t = py::ssize_t;
 
 // The maximum depth of a pytree.
-constexpr ssize_t MAX_RECURSION_DEPTH = 1000;
+#ifndef Py_C_RECURSION_LIMIT
+#define Py_C_RECURSION_LIMIT 1000
+#endif
+constexpr ssize_t MAX_RECURSION_DEPTH = std::min(1000, Py_C_RECURSION_LIMIT);
 
 // Test whether the given object is a leaf node.
 bool IsLeaf(const py::object &object,
--- include/utils.h.orig	2024-03-25 18:01:43 UTC
+++ include/utils.h
@@ -208,6 +208,9 @@ inline ssize_t GET_SIZE<py::list>(const py::handle& si
 inline ssize_t GET_SIZE<py::list>(const py::handle& sized) {
     return PyList_GET_SIZE(sized.ptr());
 }
+#ifndef PyDict_GET_SIZE
+#define PyDict_GET_SIZE PyDict_GetSize
+#endif
 template <>
 inline ssize_t GET_SIZE<py::dict>(const py::handle& sized) {
     return PyDict_GET_SIZE(sized.ptr());
@@ -222,91 +225,35 @@ template <>
     return PyTuple_GET_ITEM(container.ptr(), item);
 }
 template <>
-inline py::handle GET_ITEM_HANDLE<py::tuple>(const py::handle& container, const py::size_t& item) {
-    return PyTuple_GET_ITEM(container.ptr(), py::ssize_t_cast(item));
-}
-template <>
-inline py::handle GET_ITEM_HANDLE<py::tuple>(const py::handle& container, const int& item) {
-    return PyTuple_GET_ITEM(container.ptr(), py::ssize_t_cast(item));
-}
-template <>
 inline py::handle GET_ITEM_HANDLE<py::list>(const py::handle& container, const py::ssize_t& item) {
     return PyList_GET_ITEM(container.ptr(), item);
 }
-template <>
-inline py::handle GET_ITEM_HANDLE<py::list>(const py::handle& container, const py::size_t& item) {
-    return PyList_GET_ITEM(container.ptr(), py::ssize_t_cast(item));
-}
-template <>
-inline py::handle GET_ITEM_HANDLE<py::list>(const py::handle& container, const int& item) {
-    return PyList_GET_ITEM(container.ptr(), py::ssize_t_cast(item));
-}
 
 template <typename Container, typename Item>
 inline py::object GET_ITEM_BORROW(const py::handle& container, const Item& item) {
-    return container[item];
+    return py::reinterpret_borrow<py::object>(container[item]);
 }
 template <>
 inline py::object GET_ITEM_BORROW<py::tuple>(const py::handle& container, const py::ssize_t& item) {
     return py::reinterpret_borrow<py::object>(PyTuple_GET_ITEM(container.ptr(), item));
 }
 template <>
-inline py::object GET_ITEM_BORROW<py::tuple>(const py::handle& container, const py::size_t& item) {
-    return py::reinterpret_borrow<py::object>(
-        PyTuple_GET_ITEM(container.ptr(), py::ssize_t_cast(item)));
-}
-template <>
-inline py::object GET_ITEM_BORROW<py::tuple>(const py::handle& container, const int& item) {
-    return py::reinterpret_borrow<py::object>(
-        PyTuple_GET_ITEM(container.ptr(), py::ssize_t_cast(item)));
-}
-template <>
 inline py::object GET_ITEM_BORROW<py::list>(const py::handle& container, const py::ssize_t& item) {
     return py::reinterpret_borrow<py::object>(PyList_GET_ITEM(container.ptr(), item));
 }
-template <>
-inline py::object GET_ITEM_BORROW<py::list>(const py::handle& container, const py::size_t& item) {
-    return py::reinterpret_borrow<py::object>(
-        PyList_GET_ITEM(container.ptr(), py::ssize_t_cast(item)));
-}
-template <>
-inline py::object GET_ITEM_BORROW<py::list>(const py::handle& container, const int& item) {
-    return py::reinterpret_borrow<py::object>(
-        PyList_GET_ITEM(container.ptr(), py::ssize_t_cast(item)));
-}
 
 template <typename Container, typename Item>
 inline py::object GET_ITEM_STEAL(const py::handle& container, const Item& item) {
-    return container[item];
+    return py::reinterpret_steal<py::object>(container[item]);
 }
 template <>
 inline py::object GET_ITEM_STEAL<py::tuple>(const py::handle& container, const py::ssize_t& item) {
     return py::reinterpret_steal<py::object>(PyTuple_GET_ITEM(container.ptr(), item));
 }
 template <>
-inline py::object GET_ITEM_STEAL<py::tuple>(const py::handle& container, const py::size_t& item) {
-    return py::reinterpret_steal<py::object>(
-        PyTuple_GET_ITEM(container.ptr(), py::ssize_t_cast(item)));
-}
-template <>
-inline py::object GET_ITEM_STEAL<py::tuple>(const py::handle& container, const int& item) {
-    return py::reinterpret_steal<py::object>(
-        PyTuple_GET_ITEM(container.ptr(), py::ssize_t_cast(item)));
-}
-template <>
 inline py::object GET_ITEM_STEAL<py::list>(const py::handle& container, const py::ssize_t& item) {
     return py::reinterpret_steal<py::object>(PyList_GET_ITEM(container.ptr(), item));
 }
-template <>
-inline py::object GET_ITEM_STEAL<py::list>(const py::handle& container, const py::size_t& item) {
-    return py::reinterpret_steal<py::object>(
-        PyList_GET_ITEM(container.ptr(), py::ssize_t_cast(item)));
-}
-template <>
-inline py::object GET_ITEM_STEAL<py::list>(const py::handle& container, const int& item) {
-    return py::reinterpret_steal<py::object>(
-        PyList_GET_ITEM(container.ptr(), py::ssize_t_cast(item)));
-}
 
 template <typename Container, typename Item>
 inline void SET_ITEM(const py::handle& container, const Item& item, const py::handle& value) {
@@ -314,39 +261,15 @@ inline void SET_ITEM<py::tuple>(const py::handle& cont
 }
 template <>
 inline void SET_ITEM<py::tuple>(const py::handle& container,
-                                const ssize_t& item,
+                                const py::ssize_t& item,
                                 const py::handle& value) {
     PyTuple_SET_ITEM(container.ptr(), item, value.inc_ref().ptr());
 }
 template <>
-inline void SET_ITEM<py::tuple>(const py::handle& container,
-                                const py::size_t& item,
-                                const py::handle& value) {
-    PyTuple_SET_ITEM(container.ptr(), py::ssize_t_cast(item), value.inc_ref().ptr());
-}
-template <>
-inline void SET_ITEM<py::tuple>(const py::handle& container,
-                                const int& item,
-                                const py::handle& value) {
-    PyTuple_SET_ITEM(container.ptr(), py::ssize_t_cast(item), value.inc_ref().ptr());
-}
-template <>
 inline void SET_ITEM<py::list>(const py::handle& container,
-                               const ssize_t& item,
+                               const py::ssize_t& item,
                                const py::handle& value) {
     PyList_SET_ITEM(container.ptr(), item, value.inc_ref().ptr());
-}
-template <>
-inline void SET_ITEM<py::list>(const py::handle& container,
-                               const py::size_t& item,
-                               const py::handle& value) {
-    PyList_SET_ITEM(container.ptr(), py::ssize_t_cast(item), value.inc_ref().ptr());
-}
-template <>
-inline void SET_ITEM<py::list>(const py::handle& container,
-                               const int& item,
-                               const py::handle& value) {
-    PyList_SET_ITEM(container.ptr(), py::ssize_t_cast(item), value.inc_ref().ptr());
 }
 
 inline std::string PyRepr(const py::handle& object) {
--- src/treespec/constructor.cpp.orig	2024-03-25 18:01:43 UTC
+++ src/treespec/constructor.cpp
@@ -78,7 +78,7 @@ template <bool NoneIsLeaf>
                                             std::vector<PyTreeSpec>& treespecs,
                                             std::string& register_namespace) {
         for (const py::object& child : children) {
-            if (!py::isinstance<PyTreeSpec>(child)) {
+            if (!py::isinstance<PyTreeSpec>(child)) [[unlikely]] {
                 std::ostringstream oss{};
                 oss << "Expected a(n) " << NodeKindToString(node) << " of PyTreeSpec(s), got "
                     << PyRepr(handle) << ".";
@@ -213,15 +213,15 @@ template <bool NoneIsLeaf>
                 throw std::runtime_error(oss.str());
             }
             node.arity = 0;
-            node.node_data = GET_ITEM_BORROW<py::tuple>(out, 1);
+            node.node_data = GET_ITEM_BORROW<py::tuple>(out, ssize_t(1));
             for (const py::handle& child :
-                 py::cast<py::iterable>(GET_ITEM_BORROW<py::tuple>(out, 0))) {
+                 py::cast<py::iterable>(GET_ITEM_BORROW<py::tuple>(out, ssize_t(0)))) {
                 ++node.arity;
                 children.emplace_back(py::reinterpret_borrow<py::object>(child));
             }
             verify_children(children, treespecs, registry_namespace);
             if (num_out == 3) [[likely]] {
-                py::object node_entries = GET_ITEM_BORROW<py::tuple>(out, 2);
+                py::object node_entries = GET_ITEM_BORROW<py::tuple>(out, ssize_t(2));
                 if (!node_entries.is_none()) [[likely]] {
                     node.node_entries = py::cast<py::tuple>(std::move(node_entries));
                     const ssize_t num_entries = GET_SIZE<py::tuple>(node.node_entries);
--- src/treespec/flatten.cpp.orig	2024-03-25 18:01:43 UTC
+++ src/treespec/flatten.cpp
@@ -144,14 +144,14 @@ bool PyTreeSpec::FlattenIntoImpl(const py::handle& han
                     throw std::runtime_error(oss.str());
                 }
                 node.arity = 0;
-                node.node_data = GET_ITEM_BORROW<py::tuple>(out, 1);
+                node.node_data = GET_ITEM_BORROW<py::tuple>(out, ssize_t(1));
                 for (const py::handle& child :
-                     py::cast<py::iterable>(GET_ITEM_BORROW<py::tuple>(out, 0))) {
+                     py::cast<py::iterable>(GET_ITEM_BORROW<py::tuple>(out, ssize_t(0)))) {
                     ++node.arity;
                     recurse(child);
                 }
                 if (num_out == 3) [[likely]] {
-                    py::object node_entries = GET_ITEM_BORROW<py::tuple>(out, 2);
+                    py::object node_entries = GET_ITEM_BORROW<py::tuple>(out, ssize_t(2));
                     if (!node_entries.is_none()) [[likely]] {
                         node.node_entries = py::cast<py::tuple>(std::move(node_entries));
                         const ssize_t num_entries = GET_SIZE<py::tuple>(node.node_entries);
@@ -338,16 +338,16 @@ bool PyTreeSpec::FlattenIntoWithPathImpl(const py::han
                     throw std::runtime_error(oss.str());
                 }
                 node.arity = 0;
-                node.node_data = GET_ITEM_BORROW<py::tuple>(out, 1);
+                node.node_data = GET_ITEM_BORROW<py::tuple>(out, ssize_t(1));
                 py::object node_entries;
                 if (num_out == 3) [[likely]] {
-                    node_entries = GET_ITEM_BORROW<py::tuple>(out, 2);
+                    node_entries = GET_ITEM_BORROW<py::tuple>(out, ssize_t(2));
                 } else [[unlikely]] {
                     node_entries = py::none();
                 }
                 if (node_entries.is_none()) [[unlikely]] {
                     for (const py::handle& child :
-                         py::cast<py::iterable>(GET_ITEM_BORROW<py::tuple>(out, 0))) {
+                         py::cast<py::iterable>(GET_ITEM_BORROW<py::tuple>(out, ssize_t(0)))) {
                         recurse(child, py::int_(node.arity++));
                     }
                 } else [[likely]] {
@@ -355,7 +355,7 @@ bool PyTreeSpec::FlattenIntoWithPathImpl(const py::han
                     node.arity = GET_SIZE<py::tuple>(node.node_entries);
                     ssize_t num_children = 0;
                     for (const py::handle& child :
-                         py::cast<py::iterable>(GET_ITEM_BORROW<py::tuple>(out, 0))) {
+                         py::cast<py::iterable>(GET_ITEM_BORROW<py::tuple>(out, ssize_t(0)))) {
                         if (num_children >= node.arity) [[unlikely]] {
                             throw std::runtime_error(
                                 "PyTree custom flatten function for type " +
@@ -500,9 +500,10 @@ py::list PyTreeSpec::FlattenUpTo(const py::object& ful
             case PyTreeKind::DefaultDict: {
                 AssertExactStandardDict(object);
                 auto dict = py::reinterpret_borrow<py::dict>(object);
-                py::list expected_keys = (node.kind != PyTreeKind::DefaultDict
-                                              ? node.node_data
-                                              : GET_ITEM_BORROW<py::tuple>(node.node_data, 1));
+                py::list expected_keys =
+                    (node.kind != PyTreeKind::DefaultDict
+                         ? node.node_data
+                         : GET_ITEM_BORROW<py::tuple>(node.node_data, ssize_t(1)));
                 if (!DictKeysEqual(expected_keys, dict)) [[unlikely]] {
                     py::list keys = SortedDictKeys(dict);
                     auto [missing_keys, extra_keys] = DictKeysDifference(expected_keys, dict);
@@ -619,7 +620,7 @@ py::list PyTreeSpec::FlattenUpTo(const py::object& ful
                         << " should return a 2- or 3-tuple, got " << num_out << ".";
                     throw std::runtime_error(oss.str());
                 }
-                py::object node_data = GET_ITEM_BORROW<py::tuple>(out, 1);
+                py::object node_data = GET_ITEM_BORROW<py::tuple>(out, ssize_t(1));
                 if (node.node_data.not_equal(node_data)) [[unlikely]] {
                     std::ostringstream oss{};
                     oss << "Mismatch custom node data; expected: " << PyRepr(node.node_data)
@@ -628,7 +629,7 @@ py::list PyTreeSpec::FlattenUpTo(const py::object& ful
                 }
                 ssize_t arity = 0;
                 for (const py::handle& child :
-                     py::cast<py::iterable>(GET_ITEM_BORROW<py::tuple>(out, 0))) {
+                     py::cast<py::iterable>(GET_ITEM_BORROW<py::tuple>(out, ssize_t(0)))) {
                     ++arity;
                     agenda.emplace_back(py::reinterpret_borrow<py::object>(child));
                 }
--- src/treespec/traversal.cpp.orig	2024-03-25 18:01:43 UTC
+++ src/treespec/traversal.cpp
@@ -123,10 +123,10 @@ py::object PyTreeIter::NextImpl() {
                         << " should return a 2- or 3-tuple, got " << num_out << ".";
                     throw std::runtime_error(oss.str());
                 }
-                auto children = py::cast<py::tuple>(GET_ITEM_BORROW<py::tuple>(out, 0));
+                auto children = py::cast<py::tuple>(GET_ITEM_BORROW<py::tuple>(out, ssize_t(0)));
                 ssize_t arity = GET_SIZE<py::tuple>(children);
                 if (num_out == 3) [[likely]] {
-                    py::object node_entries = GET_ITEM_BORROW<py::tuple>(out, 2);
+                    py::object node_entries = GET_ITEM_BORROW<py::tuple>(out, ssize_t(2));
                     if (!node_entries.is_none()) [[likely]] {
                         const ssize_t num_entries =
                             GET_SIZE<py::tuple>(py::cast<py::tuple>(std::move(node_entries)));
--- src/treespec/treespec.cpp.orig	2024-03-25 18:01:43 UTC
+++ src/treespec/treespec.cpp
@@ -138,8 +138,8 @@ namespace optree {
 
         case PyTreeKind::DefaultDict: {
             py::dict dict{};
-            py::object default_factory = GET_ITEM_BORROW<py::tuple>(node.node_data, 0);
-            py::list keys = GET_ITEM_BORROW<py::tuple>(node.node_data, 1);
+            py::object default_factory = GET_ITEM_BORROW<py::tuple>(node.node_data, ssize_t(0));
+            py::list keys = GET_ITEM_BORROW<py::tuple>(node.node_data, ssize_t(1));
             if (node.original_keys) [[unlikely]] {
                 for (ssize_t i = 0; i < node.arity; ++i) {
                     dict[GET_ITEM_HANDLE<py::list>(node.original_keys, i)] = py::none();
@@ -254,11 +254,11 @@ namespace optree {
             auto expected_keys = py::reinterpret_borrow<py::list>(
                 root.kind != PyTreeKind::DefaultDict
                     ? root.node_data
-                    : GET_ITEM_BORROW<py::tuple>(root.node_data, 1));
+                    : GET_ITEM_BORROW<py::tuple>(root.node_data, ssize_t(1)));
             auto other_keys = py::reinterpret_borrow<py::list>(
                 other_root.kind != PyTreeKind::DefaultDict
                     ? other_root.node_data
-                    : GET_ITEM_BORROW<py::tuple>(other_root.node_data, 1));
+                    : GET_ITEM_BORROW<py::tuple>(other_root.node_data, ssize_t(1)));
             py::dict dict{};
             for (ssize_t i = 0; i < other_root.arity; ++i) {
                 dict[GET_ITEM_HANDLE<py::list>(other_keys, i)] = py::int_(i);
@@ -530,7 +530,7 @@ ssize_t PyTreeSpec::PathsImpl(Span& paths,
                 auto keys = py::reinterpret_borrow<py::list>(
                     root.kind != PyTreeKind::DefaultDict
                         ? root.node_data
-                        : GET_ITEM_BORROW<py::tuple>(root.node_data, 1));
+                        : GET_ITEM_BORROW<py::tuple>(root.node_data, ssize_t(1)));
                 for (ssize_t i = root.arity - 1; i >= 0; --i) {
                     cur -= recurse(cur, GET_ITEM_HANDLE<py::list>(keys, i));
                 }
@@ -595,7 +595,8 @@ py::list PyTreeSpec::Entries() const {
             return py::getattr(root.node_data, Py_Get_ID(copy))();
         }
         case PyTreeKind::DefaultDict: {
-            return py::getattr(GET_ITEM_BORROW<py::tuple>(root.node_data, 1), Py_Get_ID(copy))();
+            return py::getattr(GET_ITEM_BORROW<py::tuple>(root.node_data, ssize_t(1)),
+                               Py_Get_ID(copy))();
         }
 
         default:
@@ -631,7 +632,8 @@ py::object PyTreeSpec::Entry(ssize_t index) const {
             return GET_ITEM_BORROW<py::list>(root.node_data, index);
         }
         case PyTreeKind::DefaultDict: {
-            return GET_ITEM_BORROW<py::list>(GET_ITEM_BORROW<py::tuple>(root.node_data, 1), index);
+            return GET_ITEM_BORROW<py::list>(GET_ITEM_BORROW<py::tuple>(root.node_data, ssize_t(1)),
+                                             index);
         }
 
         case PyTreeKind::None:
@@ -804,11 +806,11 @@ bool PyTreeSpec::IsPrefix(const PyTreeSpec& other, con
                 auto expected_keys = py::reinterpret_borrow<py::list>(
                     a->kind != PyTreeKind::DefaultDict
                         ? a->node_data
-                        : GET_ITEM_BORROW<py::tuple>(a->node_data, 1));
+                        : GET_ITEM_BORROW<py::tuple>(a->node_data, ssize_t(1)));
                 auto other_keys = py::reinterpret_borrow<py::list>(
                     b->kind != PyTreeKind::DefaultDict
                         ? b->node_data
-                        : GET_ITEM_BORROW<py::tuple>(b->node_data, 1));
+                        : GET_ITEM_BORROW<py::tuple>(b->node_data, ssize_t(1)));
                 py::dict dict{};
                 for (ssize_t i = 0; i < b->arity; ++i) {
                     dict[GET_ITEM_HANDLE<py::list>(other_keys, i)] = py::int_(i);
@@ -1023,9 +1025,9 @@ std::string PyTreeSpec::ToStringImpl() const {
             case PyTreeKind::DefaultDict: {
                 EXPECT_EQ(
                     GET_SIZE<py::tuple>(node.node_data), 2, "Number of auxiliary data mismatch.");
-                py::object default_factory = GET_ITEM_BORROW<py::tuple>(node.node_data, 0);
-                auto keys =
-                    py::reinterpret_borrow<py::list>(GET_ITEM_BORROW<py::tuple>(node.node_data, 1));
+                py::object default_factory = GET_ITEM_BORROW<py::tuple>(node.node_data, ssize_t(0));
+                auto keys = py::reinterpret_borrow<py::list>(
+                    GET_ITEM_BORROW<py::tuple>(node.node_data, ssize_t(1)));
                 EXPECT_EQ(GET_SIZE<py::list>(keys),
                           node.arity,
                           "Number of keys and entries does not match.");
@@ -1146,13 +1148,13 @@ std::string PyTreeSpec::ToString() const {
             if (node.kind == PyTreeKind::DefaultDict) [[unlikely]] {
                 EXPECT_EQ(
                     GET_SIZE<py::tuple>(node.node_data), 2, "Number of auxiliary data mismatch.");
-                py::object default_factory = GET_ITEM_BORROW<py::tuple>(node.node_data, 0);
+                py::object default_factory = GET_ITEM_BORROW<py::tuple>(node.node_data, ssize_t(0));
                 data_hash = py::hash(default_factory);
             }
             auto keys = py::reinterpret_borrow<py::list>(
                 node.kind != PyTreeKind::DefaultDict
                     ? node.node_data
-                    : GET_ITEM_BORROW<py::tuple>(node.node_data, 1));
+                    : GET_ITEM_BORROW<py::tuple>(node.node_data, ssize_t(1)));
             EXPECT_EQ(
                 GET_SIZE<py::list>(keys), node.arity, "Number of keys and entries does not match.");
             for (const py::handle& key : keys) {
