Update qml.math.norm for dispatching autograd; stop rounding when postselecting

[mudit2812]

Context:
Differentiation of qml.math.norm does not work for L2 norm. This was causing incorrect gradients with autograd. Moreover, due to the rounding function, the gradient of QNodes being postselected is incorrect, so that needs to be removed.

Description of the Change:

• Add private function to compute the norm for autograd interface with qml.math.norm when ord=None and axis=None. Otherwise, we dispatch to scipy.linalg.norm as we did before.
• Stop rounding the norm of the state when renormalizing the state vector after postselection. Instead, we check if the norm is close to 0 and set it to exactly 0 if it is. This condition is only checked is the state vector is not abstract.
• Added warning to qml.measure docs about jitting with postselection on zero probability states.

Benefits:

• qml.math.norm is differentiable for all interfaces for ord=None and axis=None.
• Postselection doesn't lead to incorrect gradients.

Possible Drawbacks:
Postselection with jitting can lead to incorrect results and errors if postselecting on a state with zero probability. However, this is an edge case that is not causing problems frequently.

Related GitHub Issues:
#4867

[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.

[mudit2812]

[sc-50617]

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Comparison is base (c7cda37) 99.65% compared to head (20a982a) 99.65%.

Additional details and impacted files

@@            Coverage Diff             @@
##           master    #4766      +/-   ##
==========================================
- Coverage   99.65%   99.65%   -0.01%     
==========================================
  Files         387      387              
  Lines       34967    34711     -256     
==========================================
- Hits        34847    34590     -257     
- Misses        120      121       +1     

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

[timmysilv]

looks good! just one concern about docs.

also curious, but this might be a more general discussion - is qml.math.norm a standard thing to use? I figured we'd generally try to match the numpy API, in which case we should use qml.math.linalg.norm.

Also, if that's the only problem, I wonder if we're better off fixing autograd norms in the pennylane/numpy module instead of the math module. do you have thoughts on any of those things?

[mudit2812]

also curious, but this might be a more general discussion - is qml.math.norm a standard thing to use? I figured we'd generally try to match the numpy API, in which case we should use qml.math.linalg.norm.

There aren't any functions in qml.math that use submodules similar to the numpy API, or at least not functions implemented in multi_dispatch.py. Nothing is really stopping us from doing qml.math.linalg.norm. It works. But only with autograd and numpy, because we haven't overriden it so it automatically dispatches to <interface>.linalg.norm.

Also, if that's the only problem, I wonder if we're better off fixing autograd norms in the pennylane/numpy module instead of the math module. do you have thoughts on any of those things?

That's a good point, we could do that. @trbromley what are your thoughts about this? We could fix pennylane.numpy.linalg.norm and then dispatch to pennylane.numpy in qml.math instead of dispatching directly to autograd.numpy. The only issue I can think of is the L-inf norm, which is used in BlockEncode and is not differentiable with autograd. We could also override that 🤔

[timmysilv]

looks great! we can ignore my norm rant for now. the function already exists in multi_dispatch, and we might as well lean on that. but fwiw, we do use qml.math.linalg.matrix_power in a few places 😁