From afe94fd9f6adf68e9ee9f5b0f4cec42fb2fc64af Mon Sep 17 00:00:00 2001 From: goulustis Date: Tue, 20 Jan 2026 17:53:32 +0800 Subject: [PATCH] bug fix --- lang_agent/utils.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/lang_agent/utils.py b/lang_agent/utils.py index 3f4cce0..87f65f8 100644 --- a/lang_agent/utils.py +++ b/lang_agent/utils.py @@ -32,7 +32,9 @@ def tree_leaves(tree): while stack: node = stack.pop() if isinstance(node, dict): - stack.extend(reversed(node.values())) + # JAX sorts dict keys alphabetically + sorted_values = [node[k] for k in sorted(node.keys())] + stack.extend(reversed(sorted_values)) elif isinstance(node, (list, tuple)): stack.extend(reversed(node)) else: