[Frontend] Corrected output of PyTree when using qml.counts()

[abhamra]

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 to qml.counts with tree_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 the PyTreeDefs and correctly updates based on this traversal and a num_counts variable which stores the number of qml.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.

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 97.96%. Comparing base (af9f919) to head (af405f5).
Report is 1 commits behind head on main.

Additional details and impacted files

@@           Coverage Diff           @@
##             main    #1219   +/-   ##
=======================================
  Coverage   97.96%   97.96%           
=======================================
  Files          77       77           
  Lines       11263    11270    +7     
  Branches      972      973    +1     
=======================================
+ Hits        11034    11041    +7     
  Misses        180      180           
  Partials       49       49           

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

[mehrdad2m]

Hi @abhamra, Thanks for your submission! Your solution looks great 💯. I have added some small suggestions and comments.

Also, most PRs require a changelog entry before merging. So please add a brief description of the PR to changelog and add a url back to this PR and add your name to the list of contributors.

[abhamra]

@mehrdad2m Thank you for the feedback, I think I've addressed everything but please let me know if there's anything else. For the MCM test, I lifted the test_mcm_method_with_count_measurement from the test_mid_circuit_measurement.py file, hopefully this is what you intended.

[paul0403]

Also, @abhamra feel free to resolve conversations that have already been implemented : )

[abhamra]

@paul0403 good call on the tests and resolving comments, will do ASAP!

[mehrdad2m]

Thanks @abhamra for addressing the previous comments.

[paul0403]

@abhamra Thanks for working on this! Looks really nice ; )

I have a question unrelated to the PR, but I would like to hear your thoughts on the following phenomenon in general:

x = [1,2,3]
y = np.array([1,2,3])

jax.tree.flatten(x)
# ([1, 2, 3], PyTreeDef([*, *, *]))

jax.tree.flatten(y)
# ([array([1, 2, 3])], PyTreeDef(*))

We see that, although both objects are philosophically "a three-dimensional vector", a numpy array is considered to be just one leaf pytree node, whereas a list of three numbers is considered to be a parent node, of list type (the surrounding [] in its PyTreeDef), with three leaves.

From a compiler point of view, why is there such a distinction? What benefits would that bring?
(No need to worry about specifically how jax does whatever it does, just a discussion at an abstract level is fine)

Note that I don't have a "correct" answer that I am expecting; I too have wondered about this question at some point, so just want to hear your thoughts!

[abhamra]

@paul0403 that's a very interesting question; I had to actually verify for myself and you're right. It also seems to maintain this property with nested ndarrays as well, like so: z = np.array([[1, 2], [3, 4]]) and a = np.array([np.array([1, 2]), np.array([3, 4])]) (which folds into z) both are counted as one leaf PyTreeDef.

As for benefits, the most obvious one to me is that maintaining the structure of a numpy array allows us to exploit it later on for future optimizations (vectorization, perhaps). Operations frequently apply to the numpy array as a whole, and not to individual elements, so it makes sense to keep them together for that reason.

What's your perspective on this discrepancy?

[paul0403]

@paul0403 that's a very interesting question; I had to actually verify for myself and you're right. It also seems to maintain this property with nested ndarrays as well, like so: z = np.array([[1, 2], [3, 4]]) and a = np.array([np.array([1, 2]), np.array([3, 4])]) (which folds into z) both are counted as one leaf PyTreeDef.

As for benefits, the most obvious one to me is that maintaining the structure of a numpy array allows us to exploit it later on for future optimizations (vectorization, perhaps). Operations frequently apply to the numpy array as a whole, and not to individual elements, so it makes sense to keep them together for that reason.

What's your perspective on this discrepancy?

Interesting insights!

Yes, my understanding is also that a numpy array is basically one "thing". In a typical compiler's program representation (like mlir), variables are both typed and shaped, e.g. see the mlir tensor dialect. In fact, when jit compiling a python program, variables lose their concrete values and become abstract, only retaining their types and shapes (see #1163, "context" portion for a small introduction).

On the frontend side, a numpy array variable carries the type and shape information native in the object, so it is considered to be one whole, complete variable from the compiler's perspective. Even if the type and shape were to change, because it is a custom-defined object, it is easy to maintain and provide custom reshape/retype methods to the compiler lower down in either jaxpr or mlir.

On the other hand, native python lists don't have fixed types (i.e. the three entries in the list can have different types) and shapes. It is also native to python, and thus it's hard to tie custom methods to it.

[paul0403]

Also just a note that please don't feel obligated to diligently keep rebasing to the updated main branch. Leave the housekeeping to us 😉

(Also we are in the middle of release so branching is a bit complicated right now, so no need to worry about that from your end)

[abhamra]

@paul0403 Thank you for the insight, and w.r.t the branch updating, sounds great. Good luck with the release!