Optimise with PennyLane

examples/pennylane_expval.py

The problem

Run a PennyLane device, including a gradient step. The qubitra.remote device makes a Qubitra backend a PennyLane device, and the property worth watching is batching: PennyLane’s parameter-shift rule evaluates two shifted circuits per trainable parameter, and the device sends the whole shift batch as one multi-PUB job. Every optimisation step is then one round trip for the forward value and one for the gradient batch.

Needs the extra: pip install 'qubitra-sdk[pennylane]'.

The walkthrough

The device

import pennylane as qml
from pennylane import numpy as pnp

    # Credentials resolve from QUBITRA_API_KEY / QUBITRA_API_URL, same as QubitraClient.
    device = qml.device("qubitra.remote", wires=2, backend_id=backend_id, shots=SHOTS)

qml.device("qubitra.remote", ...) is the whole integration: the device name selects the adapter, backend_id selects the platform backend, and credentials come from the environment.

The circuit

    @qml.qnode(device, diff_method="parameter-shift")  # type: ignore[untyped-decorator]
    def circuit(params: Any) -> Any:
        # A small variational ansatz: local rotations and one entangler. <Z0 Z1> reads
        # the correlation the CNOT creates; the optimiser below pushes it toward -1.
        qml.RY(params[0], wires=0)
        qml.RY(params[1], wires=1)
        qml.CNOT(wires=[0, 1])
        return qml.expval(qml.PauliZ(0) @ qml.PauliZ(1))

A QNode over the device, differentiated by the parameter-shift rule. The rule works on hardware and simulators alike because it only ever evaluates the circuit itself, at shifted parameter values.

Forward pass

    params = pnp.array([0.4, 0.6], requires_grad=True)

    # Forward pass: one estimator PUB — the circuit without terminal measurement plus
    # ZZ as sparse Pauli terms. One job, one round trip (~seconds).
    value = circuit(params)
    print(f"<Z0 Z1> at {params.tolist()} on {backend_id}: {value:+.4f}")

Calling the QNode submits one estimator PUB and waits out the platform round trip.

The batched gradient step

    # One parameter-shift step: 2 parameters x 2 shifted evaluations = a 4-PUB batch,
    # submitted as ONE job. The wall clock of a step is a platform round trip, not
    # a round trip per shifted circuit.
    gradient = qml.grad(circuit)(params)
    params = params - STEPSIZE * gradient
    print(f"gradient: {gradient.tolist()}")
    print(f"params after one step: {params.tolist()}")

    value = circuit(params)
    print(f"<Z0 Z1> after one step: {value:+.4f} (moving toward -1)")
    return 0

qml.grad triggers the shift batch: two parameters times two shifted evaluations is four tapes, and the device submits all four as one job. The cost model that implies: a 200-iteration loop is minutes of round trips whatever the circuit costs — workable for prototyping against real backends, and the reason the batch must hold together as one job.

Running it

QUBITRA_API_KEY=qpk_... python examples/pennylane_expval.py

Output from a live run, trimmed:

<Z0 Z1> at [0.4, 0.6] on sim-statevector-26q: +0.8253
gradient: [-0.0, -0.565]
params after one step: [0.4, 0.826]
<Z0 Z1> after one step: +0.6779 (moving toward -1)

The gradient’s second component is the informative one — at these parameters only the second rotation moves the cost — and after one step of size 0.4 the expectation has moved from +0.8253 to +0.6779, toward the −1 the optimiser is pushing for.

Where to go next

• PennyLane — the full adapter: supported measurements, shot handling, and gradient batching.
• Sessions — correlating a whole optimisation run.