fix(unflatten): fix passing empty array span

metaopt · Nov 7, 2024 · 0a9dfe7 · 0a9dfe7
1 parent af33117
commit 0a9dfe7
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 1 deletion.
diff --git a/src/treespec/treespec.cpp b/src/treespec/treespec.cpp
@@ -34,6 +34,8 @@ namespace optree {
                                            const py::object children[],
                                            const size_t& num_children) {
     EXPECT_EQ(py::ssize_t_cast(num_children), node.arity, "Node arity did not match.");
+    EXPECT_TRUE(children != nullptr || num_children == 0, "Node children is null.");
+
     switch (node.kind) {
         case PyTreeKind::Leaf:
             INTERNAL_ERROR("PyTreeSpec::MakeNode() not implemented for leaves.");

diff --git a/src/treespec/unflatten.cpp b/src/treespec/unflatten.cpp
@@ -57,7 +57,9 @@ py::object PyTreeSpec::UnflattenImpl(const Span& leaves) const {
             case PyTreeKind::StructSequence:
             case PyTreeKind::Custom: {
                 const ssize_t size = py::ssize_t_cast(agenda.size());
-                py::object out = MakeNode(node, &agenda[size - node.arity], node.arity);
+                py::object out = MakeNode(node,
+                                          (node.arity > 0 ? &agenda[size - node.arity] : nullptr),
+                                          node.arity);
                 agenda.resize(size - node.arity);
                 agenda.emplace_back(std::move(out));
                 break;