fix probs.process_state with tf.function and batching #4590
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
timmysilv
commented
Sep 12, 2023
lillian542
approved these changes
Sep 12, 2023
Codecov ReportPatch coverage:
Additional details and impacted files@@ Coverage Diff @@
## master #4590 +/- ##
==========================================
- Coverage 99.65% 99.65% -0.01%
==========================================
Files 376 376
Lines 33159 33163 +4
==========================================
+ Hits 33046 33047 +1
- Misses 113 116 +3
☔ View full report in Codecov by Sentry. |
mudit2812
pushed a commit
that referenced
this pull request
Sep 12, 2023
**Context:** Tensorflow Autograph forces discipline upon me that I do not want! It doesn't like when something returns `None` sometimes and a tuple at other times. I thought I fixed `Probability.process_state` in #4470 because a test failed with `tf.function`. Unfortunately, I put the code into a state where Autograph could no longer properly compile it, so it traced all Tensors with shape `<unknown>`, and *that* made anything shape-related afterwards fail. This includes `qml.math.dot`, which `expval.process_state` does. Though I knew this was related to the problem, only yesterday was I able to single it out. And last night I realized, I ought not try to invent a fix. I should just copy from `default.qubit(.tf)`! So this PR takes code structure from [there](https://github.com/PennyLaneAI/pennylane/blob/1f93032a5f79ea0b9076d456b5b6c30515ef347d/pennylane/devices/default_qubit_tf.py#L189-L205). **Description of the Change:** Re-organize `Probability.process_state` so it allows for batching _and_ TF-Autograph support. This adds a try-expect to `ProbabilityMP._get_batch_size`, but it's what DefaultQubitTF was doing before (see link above). **Benefits:** Everything works! I am unblocked in replacing DQL with DQ2 because this broke `KerasLayer` with `model.save`! **Possible Drawbacks:** The extra try-except? I also changed the prototype for `ProbabilityMP.marginal_prob` (I removed `batch_size`) but it isn't used anywhere else
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.
Context:
Tensorflow Autograph forces discipline upon me that I do not want! It doesn't like when something returns
Nonesometimes and a tuple at other times. I thought I fixedProbability.process_statein #4470 because a test failed withtf.function. Unfortunately, I put the code into a state where Autograph could no longer properly compile it, so it traced all Tensors with shape<unknown>, and that made anything shape-related afterwards fail. This includesqml.math.dot, whichexpval.process_statedoes. Though I knew this was related to the problem, only yesterday was I able to single it out. And last night I realized, I ought not try to invent a fix. I should just copy fromdefault.qubit(.tf)! So this PR takes code structure from there.Description of the Change:
Re-organize
Probability.process_stateso it allows for batching and TF-Autograph support. This adds a try-expect toProbabilityMP._get_batch_size, but it's what DefaultQubitTF was doing before (see link above).Benefits:
Everything works! I am unblocked in replacing DQL with DQ2 because this broke
KerasLayerwithmodel.save!Possible Drawbacks:
The extra try-except? I also changed the prototype for
ProbabilityMP.marginal_prob(I removedbatch_size) but it isn't used anywhere else