Add ability to customize the tracker

[josh146]

Feature details

Currently, qml.Tracker works really well for tracking quantities that are already being recorded by the device.

However, it is not currently possible to add custom quantities to track. For example, you might want to track intermediate executions, number of executions, or even internal transforms like Hamiltonian grouping/expansion or the parameter shift rule.

Implementation

Adding custom tracking quantities can be done on the device level currently in a bit of a 'hacky' way, by subclassing the device:

import pennylane as qml
from pennylane import numpy as np
from pennylane.devices import DefaultQubit

class HackedDevice(DefaultQubit):

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.step = 0

    def execute(self, *args, **kwargs):
        res = super().execute(*args, **kwargs)
        self.step += 1

        if self.tracker.active:
            self.tracker.update(res=res, step=self.step)
            self.tracker.record()

        return res

dev = HackedDevice(wires=3, shots=1000)

@qml.qnode(dev, diff_method="parameter-shift")
def circuit(weights):
    qml.RX(weights[0], wires=0)
    qml.RX(weights[1], wires=1)
    qml.RX(weights[2], wires=2)
    qml.CNOT(wires=[0, 1])
    qml.RY(weights[4], wires=0)
    qml.RY(weights[5], wires=1)
    qml.RY(weights[6], wires=2)
    qml.CNOT(wires=[1, 2])
    qml.CRX(weights[7], wires=[0, 1])
    qml.CRY(weights[8], wires=[1, 2])
    qml.CRZ(weights[9], wires=[2, 0])
    return qml.expval(qml.PauliZ(0) @ qml.PauliZ(1))

weights = np.array([0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1.0], requires_grad=True)

def my_cb(latest, **kwargs):
    if "step" in latest:
        print(f"logging step {latest['step']}, intermediate result {latest['res']}")

with qml.Tracker(dev, callback=my_cb) as tracker:
    print(qml.grad(circuit)(weights))

In terms of the actual implementation, it would require some thinking as to the best places to make custom tracking available (devices and transforms come to mind), while also thinking about the user interface.

How important would you say this feature is?

1: Not important. Would be nice to have.

Additional information

No response

[AnuravModak]

Hello @josh146 , is this still open to work if yes kindly let me know!!

[trbromley]

Hi @AnuravModak! This is still something that could be interesting, but I'd recommend taking on a different issue as there are quite a lot of open questions surrounding how we'd like to do this. Let me know if you'd like a recommendation for which issues might be good to take on!

[AnuravModak]

Hi @AnuravModak! This is still something that could be interesting, but I'd recommend taking on a different issue as there are quite a lot of open questions surrounding how we'd like to do this. Let me know if you'd like a recommendation for which issues might be good to take on!

yeah please let me know about the issues i can work on.

[trbromley]

Hi @AnuravModak. Issue https://github.com/PennyLaneAI/pennylane/issues/2962 suggested by @albi3ro might be a good place to get started. There are more advanced issues we can suggest if you'd like to learn more. Thanks!

[AnuravModak]

sure and thanks!