catalyst
catalyst copied to clipboard
Experimental dictionary support for counts
The idea is to use a custom PyTree class that is injected into the JAXPR when we are tracing the tape. The PyTree node will remain an array tuple for as long as we tracing, and once concrete values are available it unflattens into a Python dictionary.
Benefits:
- works well and does not require a modification of the qml.counts function or a Catalyst clone thereof
Downsides:
- the problem is that qml.counts now presents two different "types" depending on whether you operate on it while tracing (pair of arrays) inside of a qjit function, or outside of a qjit function (dict)
The downside is especially problementic from a UI perspective. A full resolution might be to always present the return type of qml.counts as a dictionary, but this presents a few challenges:
- the key values are no longer "data" but "metadata", that is they are constant and cannot be produced by the program at runtime
- consequently the dictionary can also not be indexed with dynamic values (wether this is a problem or not I'm not sure)
- performance might suffer for large qubit numbers because there are 2^n dictionary elements, each its own array, vs a single array of 2^n
Another solution might involve modifying the proposed class into a user-facing data type. A type that consists of two numeric arrays but which offers a dictionary like interface:
- it prints like a dictionary (and can format the keys into bitstrings if desired)
- elements can be accessed via
__getitem__ - keys and values can be extracted as arrays
- it's immutable, but can be converted into a
dictif concrete
Hello. You may have forgotten to update the changelog!
Please edit doc/changelog.md on your branch with:
- A one-to-two sentence description of the change. You may include a small working example for new features.
- A link back to this PR.
- Your name (or GitHub username) in the contributors section.