Allow Boolean mask indexing for param_shift_hessian

[dwierichs]

Implements the option to pass a Boolean array_like as argnum to param_shift_hessian, to allow, for example, for computation of the diagonal alone. This is part of the feature request #2195 that was implemented to its major extent in #2319, but Boolean indexing was left open.
Internally, any provided argnum is converted to such a Boolean array, so that the overall code complexity is not increased as far as I can tell.

[github-actions]

Hello. You may have forgotten to update the changelog!
Please edit doc/releases/changelog-dev.md 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.

[codecov]

Codecov Report

Merging #2538 (a1ed05d) into master (6e17a6b) will increase coverage by 0.00%.
The diff coverage is 100.00%.

@@           Coverage Diff           @@
##           master    #2538   +/-   ##
=======================================
  Coverage   99.57%   99.57%           
=======================================
  Files         243      243           
  Lines       19444    19471   +27     
=======================================
+ Hits        19361    19388   +27     
  Misses         83       83           

Impacted Files	Coverage Δ
pennylane/gradients/parameter_shift_hessian.py	100.00% <100.00%> (ø)

---

Continue to review full report at Codecov.

Legend - Click here to learn more
Δ = absolute <relative> (impact), ø = not affected, ? = missing data
Powered by Codecov. Last update 6e17a6b...a1ed05d. Read the comment docs.

[dwierichs]

@cvjjm Does this work the way you would expect?

[dime10]

Thank you for adding this @dwierichs, I'll have a look at this shortly :)

[cvjjm]

Scrolling through the code this looks very good! I won't have time until next week to try it in the specific case I am interested in but I am 99% sure this does the trick.

Maybe it would be good to add one test case with params of "non-trivial shape". I.e., something like

params = np.arrray([[0.1, 0.2, 0.3], [0.4, 0.5, 0.6]])

In this case the Hessian is a 4 index tensor with shape (2, 3, 2, 3) (I think?) and thus a binary mask such as

mask = [
    [
        [[True, False, True], [True, True, False]], 
        [[False, False, True], [False, True, False]], 
        [[True, True, True], [True, False, False]]
    ], 
    [
        [[False, False, False],[False, False, False]],
        [[False, False, False],[False, False, False]],
        [[False, False, True], [True, True, True]]
    ]
]

should be possible.

If the current implementation does not support this, this can probably be added rather easily by raveling the params and mask and then, in the end, reshaping the hessian back to the proper shape.

[dwierichs]

Great @cvjjm! Regarding higher-dimensional masks, note that currently the argnum keyword argument always refers to tape parameters, and not to QNode parameters. That is, while applying param_shift_hessian to a QNode returns a Hessian of the expected shape ((qnode_outputshape, parshape, parshape)), the argnum argument is currently only applied at the tape level, so that it should always be 2d. This is certainly a discrepancy that would be good to remove as soon as possible, but currently there might not be time yet to do it.

If the current implementation does not support this, this can probably be added rather easily by raveling the params and mask > and then, in the end, reshaping the hessian back to the proper shape.

Note that in general the difference between QNode arguments and tape arguments is not only a reshaping operation or permutation but any function. In particular, the classical Jacobian needs to be considered to implement this for general QNodes.

On your end, however, it should be possible to construct intermediate tapes/QNodes that have flat parameters and thus a 2D mask anyways.

@dime10 @josh146 what do you think about the priority of allowing argnum to act on the level of QNode arguments?

[dime10]

what do you think about the priority of allowing argnum to act on the level of QNode arguments?

@dwierichs This is definitely a point of confusion I've thought about before, in that argnum for the gradient transforms refers to trainable tape parameters only (that is it subsets the set of trainable tape parameters). argnum actually setting the trainable tape parameters would make more sense to me, either directly when the input is a tape, or indirectly by referring to QNode arguments from which the set of trainable tape parameters is then defined.

[dime10]

Thank you @dwierichs great work 💯

[dime10]

Thanks @dwierichs, looks good 🎉