[Capture] Add a QmlPrimitive class to differentiate between different types of primitives

[mudit2812]

This PR adds a QmlPrimitive subclass of jax.core.Primitive. This class contains a prim_type property set using a new PrimitiveType enum. PrimitiveTypes currently available are "default", "operator", "measurement", "transform", and "higher_order". This can be made more or less fine grained as needed, but should be enough to differentiate between different types of primitives for now. Additionally, this PR:

• updates NonInterpPrimitive to be a subclass of QmlPrimitive
• updates all existing PennyLane primitives to be either QmlPrimitive or NonInterpPrimitive. See this comment to see the logic used to determine which Primitive subclass is used for each primitive.
• updates PlxprInterpreter.eval and CancelInversesInterpreter.eval to use this prim_type property.

[sc-82420]

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 99.60%. Comparing base (397273b) to head (828ad46).

Additional details and impacted files

@@           Coverage Diff           @@
##           master    #6847   +/-   ##
=======================================
  Coverage   99.60%   99.60%           
=======================================
  Files         476      477    +1     
  Lines       45182    45217   +35     
=======================================
+ Hits        45002    45037   +35     
  Misses        180      180           

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

[mudit2812]

Subclasses of PlxprInterpreter still need to be updated. Done

[albi3ro]

We could also consider adding it to the __init__?

def __init__(self, name, prim_type):
    self._prim_type = prim_type
    super().__init__(name)

[mudit2812]

We could also consider adding it to the __init__?

def __init__(self, name, prim_type):
    self._prim_type = prim_type
    super().__init__(name)

Good point. I personally like the idea of preserving the same signature as jax.core.Primitive. It's also consistent with how we currently set primitive.multiple_results. Happy to hear more thoughts and update the constructor as needed.

[PietropaoloFrisoni]

Amazing improvement 🚀

Just left a couple of totally optional suggestions

[PietropaoloFrisoni]

Suggested change

	def prim_type(self, value):
	def prim_type(self, value: Union[str, PrimitiveType]):

Just to emphasize that currently both of these work:

from pennylane.capture.custom_primitives import PrimitiveType, QmlPrimitive

prim = QmlPrimitive("primitive")

prim.prim_type = "default"

prim.prim_type = PrimitiveType("default")

in the future, it would also be nice to raise a more informative error and list all the available options, but right now this seems perfectly fine to me

[PietropaoloFrisoni]

Suggested change

	from typing import Union

[PietropaoloFrisoni]

Nice replacement. There are also similar possible replacements in tests.capture.test_operators.py (although they are just tests, so probably not relevant)