[Feature] Allow parameter shifting on all parametric gates #31
Conversation
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
atiyo
requested review from
gvelikova,
RolandMacDoland,
jpmoutinho and
chMoussa
There was a problem hiding this comment.
Thanks for this @atiyo. Not much to say, probably due to my current lack of knowledge in Jax. One thing that I would suggest is to align the syntax with PyQ. Either bu suggesting changes there or by importing them here owing the idiosyncrasies between torch and Jax ofc.
| def parse_dict(values: dict[str, float] = dict()) -> float: | ||
| return values[self.param] # type: ignore[index] | ||
| def parse_dict(self: Parametric, values: dict[str, float] = dict()) -> float: | ||
| return values[self.param_name] + self.shift # type: ignore[index] |
There was a problem hiding this comment.
Would it be possible to remove these type ignores ?
Co-authored-by: RolandMacDoland <[email protected]>
jpmoutinho
approved these changes
Oct 29, 2024
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
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.
The previous PR (#27) implements the parameter shift rule (PSR) for parameters defined in the
valuesargument of expectations. However, it suffered from some limitations:values, and a user couldn't pass a parameter directly into a gate.This MR addresses the two above points. It also adds tests that the above can be jit-compiled and give the correct answers.
Some noteworthy points:
checkify.checkpoints, the output type of the original function is changed to(error, output_of_original_function). This is not ideal for end users, so this has been removed. It would be great to have such checks in the code, so an issue investigating a promising alternative has been raised ([Enhancement] Investigate/use Chex for assertions and type handling #30).paramattribute of aParametricgate could be of typestr | float. This was problematic when implementing custom JVP rules, since afloatis a valid jax type, but a string is not (e.g. Allow str to be a valid JAX type jax-ml/jax#3045). Consequently,paramhas been explicitly split intoparam_name: strandparam_val: float, so thatparam_valis always a valid jax type.Closes #29