[Frontend] Corrected output of PyTree when using qml.counts() #1219
+131
−10
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Context: When using
qml.counts()
in the output of a quantum circuit with qjit, the output pytree is modified to replace the output pytree element related toqml.counts
withtree_structure(("keys", "counts"))
. Previously, this operation was buggy and didn't work well with nested return statements, more complex patterns, etc. This PR aims to fix this problem.Description of the Change: We add a
replace_child_tree(tree, index, subtree)
method as suggested here, which recursively traverses thePyTreeDef
s and correctly updates based on this traversal and anum_counts
variable which stores the number ofqml.counts()
calls within a potentially complex return expression.Benefits: This function should handle more complex, nested cases robustly.
Possible Drawbacks: More computationally intensive than the prior version.
Related GitHub Issues: This closes #1016.