Refactor cache_execute() as a transform

[Mandrenkov]

Context:

Presently, the cache_execute() function is difficult to understand because it assumes too many responsibilities. This PR offers an alternative implementation, _cache_transform(), which leverages the @transform infrastructure to provide the same functionality but in a shorter, cleaner, and more concise way.

If this PR is approved, we should consider deprecating cache_execute() in the next release.

Description of the Change:

• Added the following functions:
  • _cache_transform() which re-implements cache_execute() using @transform.
  • _apply_cache_transform() which conveniently wraps _cache_transform().
• Updated the existing cache_execute() tests to spy on _cache_transform() instead.
• Added a new set of tests to explicitly test the behaviour of _cache_transform().

Benefits:

• A significant developer productivity boost for those debugging cache_execute().

Possible Drawbacks:

• Some additional work will be required to deprecate and eventually remove cache_execute().

Related GitHub Issues:

None.

[Mandrenkov]

@albi3ro I've tagged you since you are the primary code owner but I think we should get someone else involved too.

[Mandrenkov]

Thanks for the comment, @albi3ro!

[albi3ro]

Sorry about forgetting this aspect.

We should also include a warning if we are using finite-shots, a persistent cache, and re-use a result:

pennylane/pennylane/workflow/execution.py

Line 385 in ff533d5

if tape.shots and getattr(cache, "_persistent_cache", True):

[Mandrenkov]

Thanks for the updated review, @albi3ro! 🥳

[albi3ro]

OK I think I see what the failures are. cache=None should be a valid kwarg to _cache_transform, in which case we just return (tape,) , null_postprocessing early.

[Mandrenkov]

OK I think I see what the failures are. cache=None should be a valid kwarg to _cache_transform, in which case we just return (tape,) , null_postprocessing early.

Ah, sorry @albi3ro, I missed that the tests failed!

In my opinion, we shouldn't support passing None to _cache_transform because it can hide bugs and encourage bad programming practices: if you don't have a cache, you shouldn't be calling _cache_transform in the first place.

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

❗ No coverage uploaded for pull request base (master@485db04). Click here to learn what that means.

Additional details and impacted files

@@            Coverage Diff            @@
##             master    #5318   +/-   ##
=========================================
  Coverage          ?   99.66%           
=========================================
  Files             ?      401           
  Lines             ?    36927           
  Branches          ?        0           
=========================================
  Hits              ?    36802           
  Misses            ?      125           
  Partials          ?        0           

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

[albi3ro]

So the code cov failures are strictly due the lack of unit tests for cache_execute. Given we are more or less marking that function for removal, I'm fine with leaving those lines uncovered. Though this would potentially be another piece of evidence for removing cache_execute immediately.

But the new transform is so much easier to follow than cache_execute 🎉 Thank you so much for handling this improvement :)

[Mandrenkov]

Thanks for the many reviews and advice, @albi3ro!

[Mandrenkov]

After offline discussion, it was decided to remove cache_execute() altogether.

[PietropaoloFrisoni]

Thanks, @Mandrenkov. This looks fantastic!
I would only suggest merging from master after 8 days to ensure everything is OK (except for the issue with lightning, which is most probably independent of this PR) 👍

[albi3ro]

Reapproving after the removal after removal of cache_execute.

[PietropaoloFrisoni]

Amazing! Thanks again @Mandrenkov