mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-19 05:16:06 +00:00
Code Issues Packages Projects Releases Wiki Activity

Avoid out-of-bounds dereference for arity-0 nodes. (#1713)

Browse Source
This commit is contained in:
Peter Hawkins 2019-11-18 15:35:07 -05:00 committed by GitHub
parent 42dd736afd
commit 9679a87901
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 7 additions and 5 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

12
jaxlib/pytree.cc
View File

@ -343,11 +343,13 @@ py::object PyTreeDef::Unflatten(py::iterable leaves) const {
case Kind::kList:
case Kind::kDict:
case Kind::kCustom: {
int size = agenda.size();
py::object o = MakeNode(
node,
absl::Span<py::object>(&agenda[size - node.arity], node.arity));
agenda.resize(agenda.size() - node.arity);
const int size = agenda.size();
absl::Span<py::object> span;
if (node.arity > 0) {
span = absl::Span<py::object>(&agenda[size - node.arity], node.arity);
}
py::object o = MakeNode(node, span);
agenda.resize(size - node.arity);
agenda.push_back(o);
break;
}