Open mehrdad2m opened 4 weeks ago
Thanks for catching this @mehrdad2m! How involved would you say the fix is -- is it straightforward, or would it require some exploration?
Thanks for catching this @mehrdad2m! How involved would you say the fix is -- is it straightforward, or would it require some exploration?
Hi @josh146, It is pretty straightforward. The fix should be done in trace_quantum_measurements which is where the output pytree is modified. Basically the simple version of the problem is to write a function replace_child_tree(tree, i, subtree) which recieved a pytree and would replace the ith node of tree that is visited in a DFS with sub_tree. The only tricky part is working with pytrees :)
Context
when using
qml.counts()
in outputof a qunatum circuit with qjit, the out put pytree is modified to replace the output pytree element related toqml.counts
withtree_structure(("keys", "counts"))
. However this tansformation is buggy and while it works in simple case, it mistransforms it in more complex patterns.A example that works fine:
The result is as expected which is:
PyTreeDef({'1': (*, *)})
In the following, there are two patterns that result in wrong output pytree:
1.
results in:
PyTreeDef(((*, *), {'2': *}))
The expected pytree is:
PyTreeDef(({'1': (*, *)}, {'2': *}))
2.
results in:
PyTreeDef(([{'1': *}, {'2': *}], (*, *)))
The expected pytree is:
PyTreeDef(([{'1': *}, {'2': (*, *)}], {'3': *}))