Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Experimental dictionary support for counts #647

Draft
wants to merge 1 commit into
base: main
Choose a base branch
from
Draft

Experimental dictionary support for counts #647

wants to merge 1 commit into from

Conversation

dime10
Copy link
Collaborator

@dime10 dime10 commented Apr 8, 2024

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 dict if concrete

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
Copy link
Contributor

github-actions bot commented Apr 8, 2024

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.

@dime10 dime10 requested a review from vincentmr April 10, 2024 21:58
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

1 participant