Make the pipeline compatible with shots decoupling

[JerryChen97]

Context:
We want to:

1. qml.set_shots or partial(set_shots, shots=...) works (the latter is the original requirement); ideally we wish it works for all interfaces and all shots/shot vec
2. previous shots setting should still work as was, until someday we decide to deprecate (ideally we wish they don't even talk with each other)

Description of the Change:

1. transforms realm is further divided (tbd): anything applied to the qnode is a user transform; (placeholder for outer); (placeholder for inner)
2. the logic relationship between the three layers of transformers:

user_transform(
  `_setup_transform_program` would decide 
  outer_transform(
    inner_transform -> presented to `run` and do whatever defined there
  ) // outer_post_processing
) // user_post_processing

Benefits:
Further clearance and decoupling for pipeline!

Possible Drawbacks:

Related GitHub Issues:
[SC-90145]

[JerryChen97]

Some technical details about what this will look like.

Right now, our execution pipeline has two transform program components. The "outer" and the "inner". The outer is run before we hit the ML boundary, while the inner is run after. The outer always contains the user transforms. If we are using a device derivative, the outer also contains device preprocessing. If we are not using device derivatives, the inner contains device preprocessing. The inner transform program also contains caching and conversion to numpy.

We now need to break this into three transform program stages. The user, the outer device, and the inner.

Right now we resolve the execution config and setup the transform programs based on the initial tape. We need to instead run the user transforms, then resolve the execution pipeline, then setup the inner and outer transform programs.

This change will cause some minor changes in tests, but I've double checked, and they are are all acceptable differences. This will also solve some long standing bugs, and allow us to get rid of this patch:

pennylane/pennylane/workflow/resolution.py

Line 283 in 3229b05

if (

Once we update the workflow, we will also need to update construct_batch and get_transform_program. We will now always be applying final transforms at the end of the user transform program, and before device preprocessing.

[JerryChen97]

Trigger the CI to see if our program separation breaks anything unexpectedly.

[JerryChen97]

I think we still need a new set to record who belong to user transforms who does not or it might be too much to adjust to make things work as expected @albi3ro

[albi3ro]

I think we still need a new set to record who belong to user transforms who does not or it might be too much to adjust to make things work as expected @albi3ro

Mind rephrasing?

[JerryChen97]

I think we still need a new set to record who belong to user transforms who does not or it might be too much to adjust to make things work as expected @albi3ro

Mind rephrasing?

Basically only do

    ### Apply the user transforms ####
    if transform_program:
        tapes, _ = transform_program(tapes)
        transform_program = TransformProgram()

for those that we recognized as a suitable "user transform", e.g. set_shots

[JerryChen97]

Like in this commit although right now it is not actually working since TransformProgram does not know the newly added attr to set_shots.
If we just apply them in the very beginning then some transform e.g. _eigvals_transform will be executed wrongly.
To make this actually meaningful, I need to add an attr is_user_transform to TransformProgram , what do you think? @albi3ro

[JerryChen97]

It looks like not immediately removable. Mind I revert commits back to c5ca4e7 and make another story for it as follow-up of pipeline clean-ups? @albi3ro

[JerryChen97]

Looks good! Only need to adjust trhee tests inside the same file seemingly!

[JerryChen97]

@albi3ro Are we safe to remove this entirely:

def test_prune_dynamic_transform_warning_raised():
    """Tests that a warning raised when a user-applied dynamic-one-shot transform is ignored."""

    user_transform_program = TransformProgram()
    user_transform_program.add_transform(qml.transforms.dynamic_one_shot)
    device = qml.device("default.qubit")
    config = ExecutionConfig(mcm_config=MCMConfig(mcm_method="one-shot"))

    with pytest.warns(UserWarning, match="A dynamic_one_shot transform already exists"):
        _, __ = _setup_transform_program(device, config)

[albi3ro]

@albi3ro Are we safe to remove this entirely:

def test_prune_dynamic_transform_warning_raised():
    """Tests that a warning raised when a user-applied dynamic-one-shot transform is ignored."""

    user_transform_program = TransformProgram()
    user_transform_program.add_transform(qml.transforms.dynamic_one_shot)
    device = qml.device("default.qubit")
    config = ExecutionConfig(mcm_config=MCMConfig(mcm_method="one-shot"))

    with pytest.warns(UserWarning, match="A dynamic_one_shot transform already exists"):
        _, __ = _setup_transform_program(device, config)

Sorry about that. Might actually need to think about where to place that warning now. We should make sure a suitable warning is raised if someone puts dynamic_one_shot on the qnode.

[JerryChen97]

@albi3ro Are we safe to remove this entirely:

def test_prune_dynamic_transform_warning_raised():
    """Tests that a warning raised when a user-applied dynamic-one-shot transform is ignored."""

    user_transform_program = TransformProgram()
    user_transform_program.add_transform(qml.transforms.dynamic_one_shot)
    device = qml.device("default.qubit")
    config = ExecutionConfig(mcm_config=MCMConfig(mcm_method="one-shot"))

    with pytest.warns(UserWarning, match="A dynamic_one_shot transform already exists"):
        _, __ = _setup_transform_program(device, config)

Sorry about that. Might actually need to think about where to place that warning now. We should make sure a suitable warning is raised if someone puts dynamic_one_shot on the qnode.

ohoh. yeah indeed. I felt this test is special since it's not directly expected by our refactoring but seems some design from mcm/dynamic-one-shot side.