# Copyright (c) 2026, Jiun-Cheng Jiang. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""
QKAN solver for real quantum device execution via Qiskit Runtime.
"""
import math
import warnings
from typing import Any, Callable, Optional
import torch
from ._base import QKANSolver, register
try:
from qiskit import QuantumCircuit # type: ignore
from qiskit.quantum_info import SparsePauliOp # type: ignore
from qiskit.transpiler.preset_passmanagers import ( # type: ignore
generate_preset_pass_manager,
)
_QISKIT_AVAILABLE = True
except ImportError:
_QISKIT_AVAILABLE = False
try:
from qiskit_ibm_runtime import EstimatorV2 as Estimator # type: ignore
_QISKIT_RUNTIME_AVAILABLE = True
except ImportError:
_QISKIT_RUNTIME_AVAILABLE = False
try:
from qiskit.primitives import StatevectorEstimator # type: ignore
_SV_ESTIMATOR_AVAILABLE = True
except ImportError:
_SV_ESTIMATOR_AVAILABLE = False
try:
from qiskit_aer import AerSimulator # type: ignore
_AER_AVAILABLE = True
except ImportError:
_AER_AVAILABLE = False
from ._mitigation import _apply_mitigation
JobEventCallback = Callable[[dict[str, Any]], None]
class _JobEventCallbackError(RuntimeError):
"""Raised when a job event cannot be persisted by its callback."""
def _configure_estimator(
rt_estimator,
shots,
resilience_level,
twirling,
*,
max_execution_time=None,
job_tags=None,
):
"""Apply shots, resilience, and twirling options to an EstimatorV2."""
if shots is not None:
rt_estimator.options.default_shots = shots
if resilience_level is not None:
rt_estimator.options.resilience_level = resilience_level
if twirling is not None:
if twirling.get("enable_gates"):
rt_estimator.options.twirling.enable_gates = True
if twirling.get("enable_measure"):
rt_estimator.options.twirling.enable_measure = True
if twirling.get("num_randomizations") is not None:
rt_estimator.options.twirling.num_randomizations = twirling[
"num_randomizations"
]
if max_execution_time is not None:
rt_estimator.options.max_execution_time = max_execution_time
if job_tags is not None:
rt_estimator.options.environment.job_tags = job_tags
def _job_event_context(base_context, **updates):
"""Return a fresh context dict for one Runtime job event."""
context = dict(base_context or {})
context.update(updates)
return context
def _emit_job_event(
job_event_callback: JobEventCallback | None,
job,
event: str,
context: dict[str, Any] | None = None,
error: Exception | None = None,
):
"""Emit a serializable submission/finalization event for a primitive job."""
if job_event_callback is None:
return
record = _job_event_context(context, event=event, job_id=job.job_id())
if error is not None:
record["error"] = {
"type": type(error).__name__,
"message": str(error),
}
if event != "submitted":
metrics = None
try:
metrics = job.metrics()
record["metrics"] = metrics
except Exception as metadata_error:
record["metrics_error"] = {
"type": type(metadata_error).__name__,
"message": str(metadata_error),
}
usage_seconds = None
if isinstance(metrics, dict):
usage = metrics.get("usage")
if isinstance(usage, dict):
usage_seconds = usage.get("quantum_seconds")
elif metrics is None:
# job.usage() re-fetches the same metrics payload, so only fall
# back to it when the metrics fetch itself failed.
try:
usage_seconds = job.usage()
except Exception as usage_error:
record["usage_error"] = {
"type": type(usage_error).__name__,
"message": str(usage_error),
}
record["usage_seconds"] = usage_seconds
try:
job_event_callback(record)
except Exception as callback_error:
raise _JobEventCallbackError("job event callback failed") from callback_error
def _submit_job(
est,
pubs,
*,
job_event_callback: JobEventCallback | None = None,
context: dict[str, Any] | None = None,
):
"""Submit one primitive job and expose its ID before waiting for results."""
job = est.run(pubs)
_emit_job_event(job_event_callback, job, "submitted", context)
return job
def _collect_job(
job,
*,
job_event_callback: JobEventCallback | None = None,
context: dict[str, Any] | None = None,
):
"""Collect one primitive job and expose its final metrics and usage."""
try:
result = job.result()
except Exception as error:
_emit_job_event(job_event_callback, job, "failed", context, error)
raise
_emit_job_event(job_event_callback, job, "completed", context)
return result
# ---------------------------------------------------------------------------
# Qiskit circuit builders
# ---------------------------------------------------------------------------
def _fold_qiskit_circuit(qc: "QuantumCircuit", scale_factor: int) -> "QuantumCircuit":
"""Apply gate folding to a Qiskit circuit for ZNE.
Produces U . (U_dag . U)^((scale_factor-1)/2) which has the same unitary
as U but with scale_factor x the gate count (and thus noise).
"""
if scale_factor <= 1:
return qc
folded = qc.copy()
for _ in range((scale_factor - 1) // 2):
folded = folded.compose(qc.inverse()).compose(qc)
return folded
def _build_qiskit_pz_circuit(
x_val: float,
theta_vals: list[float],
reps: int,
encoded_x_vals: Optional[list[float]] = None,
) -> "QuantumCircuit":
"""
Build pz_encoding circuit: H -> [RZ(θ₀) RY(θ₁) RZ(x)]^reps -> RZ(θ_f₀) RY(θ_f₁)
theta_vals layout: [θ₀₀, θ₀₁, θ₁₀, θ₁₁, ..., θ_f₀, θ_f₁] (2 params per layer + 2 final)
"""
qc = QuantumCircuit(1)
qc.h(0)
for l in range(reps):
qc.rz(theta_vals[2 * l], 0)
qc.ry(theta_vals[2 * l + 1], 0)
enc = encoded_x_vals[l] if encoded_x_vals is not None else x_val
qc.rz(enc, 0)
qc.rz(theta_vals[2 * reps], 0)
qc.ry(theta_vals[2 * reps + 1], 0)
return qc
def _build_qiskit_rpz_circuit(
encoded_x_vals: list[float],
theta_vals: list[float],
reps: int,
) -> "QuantumCircuit":
"""
Build rpz_encoding circuit: H -> [RY(θ) RZ(encoded_x)]^reps -> RY(θ_final)
theta_vals layout: [θ₀, θ₁, ..., θ_final] (1 param per layer + 1 final)
"""
qc = QuantumCircuit(1)
qc.h(0)
for l in range(reps):
qc.ry(theta_vals[l], 0)
qc.rz(encoded_x_vals[l], 0)
qc.ry(theta_vals[reps], 0)
return qc
def _build_qiskit_real_circuit(
x_val: float,
theta_vals: list[float],
reps: int,
encoded_x_vals: Optional[list[float]] = None,
) -> "QuantumCircuit":
"""
Build real ansatz circuit: H -> [X RY(θ) Z RY(x)]^reps
theta_vals layout: [θ₀, θ₁, ...] (1 param per layer, no final gate)
"""
qc = QuantumCircuit(1)
qc.h(0)
for l in range(reps):
qc.x(0)
qc.ry(theta_vals[l], 0)
qc.z(0)
enc = encoded_x_vals[l] if encoded_x_vals is not None else x_val
qc.ry(enc, 0)
return qc
# ---------------------------------------------------------------------------
# Parallel multi-qubit packing (Qiskit)
# ---------------------------------------------------------------------------
def _build_qiskit_parallel_circuit(
single_circuits: list["QuantumCircuit"],
) -> "QuantumCircuit":
"""
Pack N independent single-qubit circuits into one N-qubit circuit.
Each single-qubit circuit is applied to a separate qubit, enabling
parallel execution on a multi-qubit QPU.
"""
n = len(single_circuits)
qc = QuantumCircuit(n)
for qubit_idx, sc in enumerate(single_circuits):
for instruction in sc.data:
gate = instruction.operation
params = gate.params
name = gate.name
if name == "h":
qc.h(qubit_idx)
elif name == "x":
qc.x(qubit_idx)
elif name == "z":
qc.z(qubit_idx)
elif name == "rx":
qc.rx(params[0], qubit_idx)
elif name == "ry":
qc.ry(params[0], qubit_idx)
elif name == "rz":
qc.rz(params[0], qubit_idx)
else:
raise ValueError(f"Unsupported gate '{name}' in parallel packing")
return qc
def _make_parallel_observables(n_qubits: int) -> list["SparsePauliOp"]:
"""
Create Z observables for each qubit in an N-qubit circuit.
Returns a list of N SparsePauliOp, each measuring Z on one qubit.
Qiskit uses little-endian ordering: qubit 0 is the rightmost character.
E.g. for 3 qubits: [IIZ, IZI, ZII] for qubits 0, 1, 2 respectively.
"""
observables = []
for k in range(n_qubits):
# Qiskit little-endian: qubit k is at string position (n-1-k) from the left
pauli_str = "I" * (n_qubits - 1 - k) + "Z" + "I" * k
observables.append(SparsePauliOp.from_list([(pauli_str, 1.0)]))
return observables
# ---------------------------------------------------------------------------
# Qiskit solver
# ---------------------------------------------------------------------------
class _QiskitParamShift(torch.autograd.Function):
"""Autograd function using parameter-shift rule for Qiskit circuits."""
@staticmethod
def forward(ctx, x, theta, preacts_w, preacts_b, reps, config):
ctx.save_for_backward(x, theta, preacts_w, preacts_b)
ctx.reps = reps
ctx.config = config
return _qiskit_evaluate(x, theta, preacts_w, preacts_b, reps, config)
@staticmethod
def backward(ctx, grad_output):
x, theta, preacts_w, preacts_b = ctx.saved_tensors
reps = ctx.reps
config = ctx.config
shift = math.pi / 2
# Gradient w.r.t. theta via parameter-shift rule
grad_theta = torch.zeros_like(theta)
flat_theta = theta.reshape(-1)
for k in range(flat_theta.numel()):
theta_plus = flat_theta.clone()
theta_plus[k] += shift
theta_minus = flat_theta.clone()
theta_minus[k] -= shift
f_plus = _qiskit_evaluate(
x, theta_plus.reshape(theta.shape), preacts_w, preacts_b, reps, config
)
f_minus = _qiskit_evaluate(
x, theta_minus.reshape(theta.shape), preacts_w, preacts_b, reps, config
)
grad_k = (f_plus - f_minus) / (2 * math.sin(shift))
grad_theta.reshape(-1)[k] = (grad_output * grad_k).sum()
# Gradient w.r.t. preacts_weight
grad_pw = None
if preacts_w.requires_grad:
grad_pw = torch.zeros_like(preacts_w)
flat_pw = preacts_w.reshape(-1)
for k in range(flat_pw.numel()):
pw_plus = flat_pw.clone()
pw_plus[k] += shift
pw_minus = flat_pw.clone()
pw_minus[k] -= shift
f_plus = _qiskit_evaluate(
x, theta, pw_plus.reshape(preacts_w.shape), preacts_b, reps, config
)
f_minus = _qiskit_evaluate(
x, theta, pw_minus.reshape(preacts_w.shape), preacts_b, reps, config
)
grad_pw.reshape(-1)[k] = (
grad_output * (f_plus - f_minus) / (2 * math.sin(shift))
).sum()
# Gradient w.r.t. preacts_bias
grad_pb = None
if preacts_b.requires_grad:
grad_pb = torch.zeros_like(preacts_b)
flat_pb = preacts_b.reshape(-1)
for k in range(flat_pb.numel()):
pb_plus = flat_pb.clone()
pb_plus[k] += shift
pb_minus = flat_pb.clone()
pb_minus[k] -= shift
f_plus = _qiskit_evaluate(
x, theta, preacts_w, pb_plus.reshape(preacts_b.shape), reps, config
)
f_minus = _qiskit_evaluate(
x, theta, preacts_w, pb_minus.reshape(preacts_b.shape), reps, config
)
grad_pb.reshape(-1)[k] = (
grad_output * (f_plus - f_minus) / (2 * math.sin(shift))
).sum()
return None, grad_theta, grad_pw, grad_pb, None, None
def _probe_max_pubs(
est,
probe_pubs,
max_pubs,
*,
job_event_callback: JobEventCallback | None = None,
job_context: dict[str, Any] | None = None,
):
"""
Binary-search for the largest PUB batch the QPU accepts.
Submits `probe_pubs[:max_pubs]` synchronously. On memory error (6073),
halves and retries until a working size is found. Returns (result, max_pubs)
where result is the successful job result for the probe batch.
"""
attempt = 1
while max_pubs >= 1:
batch = probe_pubs[:max_pubs]
context = _job_event_context(
job_context,
phase="probe",
attempt=attempt,
pub_start=0,
pub_end=max_pubs,
pub_count=len(batch),
)
try:
job = _submit_job(
est,
batch,
job_event_callback=job_event_callback,
context=context,
)
result = _collect_job(
job,
job_event_callback=job_event_callback,
context=context,
)
return result, max_pubs
except _JobEventCallbackError:
raise
except Exception as e:
err_str = str(e)
if "6073" in err_str or "memory" in err_str.lower():
old_max = max_pubs
max_pubs = max(1, max_pubs // 2)
if max_pubs == old_max:
raise # can't go smaller than 1
print(
f" [qsolver] Job memory limit hit at {old_max} PUBs/job, "
f"trying {max_pubs}"
)
attempt += 1
else:
raise
raise RuntimeError("Could not find a working PUB batch size")
def _submit_and_collect(
est,
all_pubs,
all_chunk_sizes,
max_pubs,
*,
job_event_callback: JobEventCallback | None = None,
job_context: dict[str, Any] | None = None,
):
"""
Submit PUBs with the largest batch size the QPU can handle.
1. Probes with max_pubs (all PUBs if 0) synchronously to find the
largest accepted batch size via binary search on memory errors.
2. Submits all remaining batches asynchronously for max throughput.
3. Collects results in order.
Returns (expvals, actual_max_pubs) so callers can cache the working size.
"""
n_total = len(all_pubs)
if max_pubs <= 0:
max_pubs = n_total
expvals: list[Optional[list[float]]] = [None] * n_total
# Step 1: Probe with first batch to discover working max_pubs
first_batch_size = min(max_pubs, n_total)
first_batch = all_pubs[:first_batch_size]
probe_result, max_pubs = _probe_max_pubs(
est,
first_batch,
first_batch_size,
job_event_callback=job_event_callback,
job_context=job_context,
)
# Collect probe results (first max_pubs PUBs)
probed_count = min(max_pubs, n_total)
for i in range(probed_count):
evs = probe_result[i].data.evs
expvals[i] = [float(v) for v in evs]
# Step 2: Submit remaining batches asynchronously
remaining_start = probed_count
if remaining_start < n_total:
jobs: list[Any] = []
job_ranges: list[tuple[int, int]] = []
job_contexts: list[dict[str, Any]] = []
for batch_start in range(remaining_start, n_total, max_pubs):
batch_end = min(batch_start + max_pubs, n_total)
job_pubs = all_pubs[batch_start:batch_end]
context = _job_event_context(
job_context,
phase="batch",
batch_index=len(jobs),
pub_start=batch_start,
pub_end=batch_end,
pub_count=len(job_pubs),
logical_circuit_count=sum(all_chunk_sizes[batch_start:batch_end]),
)
jobs.append(
_submit_job(
est,
job_pubs,
job_event_callback=job_event_callback,
context=context,
)
)
job_ranges.append((batch_start, batch_end))
job_contexts.append(context)
n_jobs = len(jobs)
print(f" [qsolver] Submitting {n_jobs} async job(s), {max_pubs} PUBs/job")
# Collect all async results. When accounting is enabled, finalize every
# already-submitted job before re-raising the first execution failure.
first_job_error: Optional[Exception] = None
for job, (batch_start, batch_end), context in zip(
jobs, job_ranges, job_contexts
):
try:
result = _collect_job(
job,
job_event_callback=job_event_callback,
context=context,
)
except _JobEventCallbackError:
raise
except Exception as job_error:
if job_event_callback is None:
raise
if first_job_error is None:
first_job_error = job_error
continue
for i, global_idx in enumerate(range(batch_start, batch_end)):
evs = result[i].data.evs
expvals[global_idx] = [float(v) for v in evs]
if first_job_error is not None:
raise first_job_error
# Flatten
flat: list[float] = []
for ev_list in expvals:
if ev_list is not None:
flat.extend(ev_list)
return flat, max_pubs
# Module-level cache for the discovered max PUBs per backend
_MAX_PUBS_CACHE: dict = {}
[docs]
def best_qubits(
backend,
n: int,
*,
max_readout_error: float | None = None,
qubit_error_threshold: float | None = None,
strict: bool = True,
) -> list[int]:
"""Return the ``n`` best-calibrated qubit indices on ``backend``.
When packing ``n`` independent single-qubit QKAN circuits onto ``n``
physical qubits of one multi-qubit job, the naive transpiler layout
(qubits ``0..n-1``) often includes poorly-calibrated qubits. Readout
and single-qubit gate errors on the selected qubits directly bias
every expectation value in the packed job, so a few bad qubits
inflate the aggregate error of the whole batch.
This helper scores each physical qubit by
.. math::
\\mathrm{score}(q) = \\mathrm{readout\\_error}(q) +
\\mathrm{sx\\_err}(q) +
10^{-4} / \\max(T_2(q)\\,[\\mu s],\\, 1)
(readout error dominates, :math:`sx` error is secondary, short
:math:`T_2` gets a small penalty) and returns the indices of the
``n`` lowest-scoring qubits, best first. Qubits flagged
non-operational are skipped, and calibration values reported as NaN
(typical for faulty qubits) are treated as missing data rather than
ranked. Pass the result as ``initial_layout`` via ``solver_kwargs``
to pin the packed circuit onto them.
Empirically on ``FakeSherbrooke`` with ``parallel_qubits=20``,
``shots=1024``: the smart layout recovers near-``parallel_qubits=1``
fidelity (≈0.13% rel MSE vs a noiseless reference) where the naive
layout lands at ≈5% rel MSE — a ~40× improvement at identical QPU
cost.
Parameters
----------
backend : qiskit Backend
Backend with a ``properties()`` method (FakeProvider or real IBM).
Returns an empty list if the backend exposes no calibration and
no threshold was requested.
n : int
Number of qubits to select. Must be a positive integer and must
not exceed ``backend.num_qubits``.
max_readout_error : float, optional
Keep only qubits with readout error at or below this value.
qubit_error_threshold : float, optional
Keep only qubits whose single-qubit ``sx`` gate error is at or
below this value; e.g. ``0.001`` means fidelity >= 99.9%.
strict : bool
When a threshold leaves fewer than ``n`` usable qubits, raise
``ValueError`` if true (default); otherwise return all usable
qubits.
Returns
-------
list[int]
Top-``n`` physical qubit indices, sorted by ascending score.
Raises
------
ValueError
If ``n`` is out of range, if a requested threshold cannot be
enforced because calibration is unavailable, or (with ``strict``)
if fewer than ``n`` qubits satisfy the thresholds.
Examples
--------
>>> from qiskit_ibm_runtime.fake_provider import FakeSherbrooke
>>> backend = FakeSherbrooke()
>>> layout = best_qubits(backend, 20)
>>> model = QKAN(
... [1, 2, 1], solver="qiskit", fast_measure=False,
... solver_kwargs={
... "backend": backend,
... "shots": 1024,
... "parallel_qubits": 20,
... "initial_layout": layout,
... },
... )
"""
if not isinstance(n, int) or n < 1:
raise ValueError(f"best_qubits: n must be a positive integer, got {n!r}")
has_threshold = max_readout_error is not None or qubit_error_threshold is not None
try:
props = backend.properties()
except Exception as exc:
if has_threshold:
raise ValueError(
"best_qubits: backend calibration is unavailable; "
"cannot enforce qubit calibration thresholds"
) from exc
return []
if props is None:
if has_threshold:
raise ValueError(
"best_qubits: backend calibration is unavailable; "
"cannot enforce qubit calibration thresholds"
)
return []
try:
num_qubits = backend.num_qubits
except Exception:
return []
if n > num_qubits:
raise ValueError(
f"best_qubits: requested n={n} exceeds backend.num_qubits={num_qubits}"
)
scored = []
for q in range(num_qubits):
try:
if not props.is_qubit_operational(q):
continue
except Exception:
pass
# Faulty qubits report NaN calibration values; NaN compares False
# against any threshold and corrupts sorting, so treat NaN exactly
# like missing data: it fails an active threshold and falls back to
# a pessimistic constant when only ranking.
try:
ro = float(props.readout_error(q))
except Exception:
ro = float("nan")
if math.isnan(ro):
if max_readout_error is not None:
continue
ro = 0.5
if max_readout_error is not None and ro > max_readout_error:
continue
try:
sx = float(props.gate_error("sx", [q]))
except Exception:
sx = float("nan")
if math.isnan(sx):
if qubit_error_threshold is not None:
continue
sx = 1e-2
if qubit_error_threshold is not None and sx > qubit_error_threshold:
continue
try:
t2_us = float(props.t2(q)) * 1e6
except Exception:
t2_us = float("nan")
if math.isnan(t2_us):
t2_us = 50.0
score = ro + sx + 1e-4 / max(t2_us, 1.0)
scored.append((score, q))
if len(scored) < n:
thresholds = []
if max_readout_error is not None:
thresholds.append(f"max_readout_error<={max_readout_error:.6g}")
if qubit_error_threshold is not None:
thresholds.append(f"single_qubit_sx_error<={qubit_error_threshold:.6g}")
threshold_text = " and ".join(thresholds) if thresholds else "ranking"
if not strict and scored:
scored.sort()
return [q for _score, q in scored]
raise ValueError(
f"best_qubits: only {len(scored)} qubits satisfy {threshold_text}; need {n}"
)
scored.sort()
return [q for _score, q in scored[:n]]
def _initial_layout_for_circuit(initial_layout, circuit_qubits: int):
"""Return an initial layout with the same length as the circuit.
qiskit rejects a layout whose length differs from the circuit width, so
packed chunks narrower than ``parallel_qubits`` (ragged final chunk,
or a total circuit count below the packing width) use the first
``circuit_qubits`` entries — the best-ranked qubits when the layout
came from :func:`best_qubits`.
"""
if initial_layout is None:
return None
if isinstance(initial_layout, tuple):
initial_layout = list(initial_layout)
elif hasattr(initial_layout, "tolist"):
# numpy arrays / torch tensors would otherwise skip the truncation
# below and qiskit rejects layouts longer than the circuit.
initial_layout = list(initial_layout.tolist())
if isinstance(initial_layout, list):
if not initial_layout:
# qiskit treats an empty layout like None; keep that fallback.
return None
if len(initial_layout) < circuit_qubits:
raise ValueError(
"initial_layout has fewer qubits than the circuit: "
f"{len(initial_layout)} < {circuit_qubits}"
)
return initial_layout[:circuit_qubits]
return initial_layout
def _qiskit_run_parallel(
circuits,
n_qubits,
estimator,
backend,
optimization_level,
shots,
initial_layout=None,
max_pubs_per_job=0,
resilience_level=None,
twirling=None,
max_execution_time=None,
job_tags=None,
job_event_callback: JobEventCallback | None = None,
job_context: dict[str, Any] | None = None,
):
"""
Pack single-qubit circuits into multi-qubit batches and submit async.
Groups `circuits` into chunks of `n_qubits`, packs each chunk into one
multi-qubit circuit. Jobs are submitted asynchronously with automatic
PUB batch sizing:
- If `max_pubs_per_job` > 0, uses that as the initial batch size.
- If `max_pubs_per_job` == 0 (default), starts with all PUBs in one job.
- On memory error (6073), automatically halves and retries.
- The discovered working batch size is cached per backend.
When `initial_layout` is a list of `n_qubits` physical qubit indices,
the packed circuit is pinned to those qubits during transpilation.
The layout is truncated to each chunk's width, so a ragged final
chunk uses the first `chunk_size` entries — the best-ranked qubits
when the layout came from :func:`best_qubits`. This is how
:func:`best_qubits` gets applied — qkan does not auto-select qubits
unless the caller explicitly requests it.
"""
total = len(circuits)
# Build all PUBs first
all_pubs = []
all_chunk_sizes = []
if estimator is not None:
for start in range(0, total, n_qubits):
end = min(start + n_qubits, total)
batch_circuits = circuits[start:end]
chunk_size = end - start
all_chunk_sizes.append(chunk_size)
packed_qc = _build_qiskit_parallel_circuit(batch_circuits)
chunk_obs = _make_parallel_observables(chunk_size)
all_pubs.append((packed_qc, chunk_obs))
initial_max = max_pubs_per_job if max_pubs_per_job > 0 else len(all_pubs)
expvals, _ = _submit_and_collect(
estimator,
all_pubs,
all_chunk_sizes,
initial_max,
job_event_callback=job_event_callback,
job_context=job_context,
)
return expvals
elif backend is not None:
# One pass manager per chunk width: the final chunk can be narrower
# than n_qubits, and qiskit pins the layout length at construction.
pm_cache: dict = {}
rt_estimator = Estimator(mode=backend)
_configure_estimator(
rt_estimator,
shots,
resilience_level,
twirling,
max_execution_time=max_execution_time,
job_tags=job_tags,
)
for start in range(0, total, n_qubits):
end = min(start + n_qubits, total)
batch_circuits = circuits[start:end]
chunk_size = end - start
all_chunk_sizes.append(chunk_size)
packed_qc = _build_qiskit_parallel_circuit(batch_circuits)
if chunk_size not in pm_cache:
pm_cache[chunk_size] = generate_preset_pass_manager(
backend=backend,
optimization_level=optimization_level,
initial_layout=_initial_layout_for_circuit(
initial_layout, chunk_size
),
)
isa_qc = pm_cache[chunk_size].run(packed_qc)
chunk_obs = _make_parallel_observables(chunk_size)
isa_obs = [obs.apply_layout(isa_qc.layout) for obs in chunk_obs]
all_pubs.append((isa_qc, isa_obs))
# Use cached max or start with all PUBs
cache_key = getattr(backend, "name", str(backend))
initial_max = (
max_pubs_per_job
if max_pubs_per_job > 0
else _MAX_PUBS_CACHE.get(cache_key, len(all_pubs))
)
expvals, actual_max = _submit_and_collect(
rt_estimator,
all_pubs,
all_chunk_sizes,
initial_max,
job_event_callback=job_event_callback,
job_context=job_context,
)
_MAX_PUBS_CACHE[cache_key] = actual_max
return expvals
return []
def _qiskit_evaluate(
x: torch.Tensor,
theta: torch.Tensor,
preacts_weight: torch.Tensor,
preacts_bias: torch.Tensor,
reps: int,
config: dict,
) -> torch.Tensor:
"""
Evaluate all circuits on the Qiskit backend and return expectation values.
Returns shape: (batch_size, out_dim, in_dim)
"""
batch, in_dim = x.shape
ansatz = config["ansatz"]
preacts_trainable = config["preacts_trainable"]
out_dim = config["out_dim"]
backend = config.get("backend", None)
estimator = config.get("estimator", None)
shots = config["shots"]
optimization_level = config.get("optimization_level", 1)
parallel_qubits = config.get("parallel_qubits", None)
initial_layout = config.get("initial_layout", None)
# Broadcast theta/preacts to (out_dim, in_dim, ...)
if len(theta.shape) != 4:
theta = theta.unsqueeze(0)
if theta.shape[1] != in_dim:
repeat_out = out_dim
repeat_in = in_dim // theta.shape[1] + 1
theta = theta.repeat(repeat_out, repeat_in, 1, 1)[:, :in_dim, :, :]
_needs_encoded_x = preacts_trainable or ansatz in ("rpz_encoding", "rpz")
encoded_x = None
if _needs_encoded_x:
if len(preacts_weight.shape) != 3:
preacts_weight = preacts_weight.unsqueeze(0)
preacts_bias = preacts_bias.unsqueeze(0)
if preacts_weight.shape[1] != in_dim:
repeat_out = out_dim
repeat_in = in_dim // preacts_weight.shape[1] + 1
preacts_weight = preacts_weight.repeat(repeat_out, repeat_in, 1)[
:, :in_dim, :
]
preacts_bias = preacts_bias.repeat(repeat_out, repeat_in, 1)[:, :in_dim, :]
encoded_x = torch.einsum("oir,bi->boir", preacts_weight, x).add(preacts_bias)
# Move to CPU for circuit parameter extraction
x_np = x.detach().cpu()
theta_np = theta.detach().cpu()
encoded_x_np = encoded_x.detach().cpu() if encoded_x is not None else None
# Build circuits and observables
circuits = []
observables = []
pauli_z = SparsePauliOp.from_list([("Z", 1.0)])
# Pre-convert theta (batch-independent) and x to Python lists
theta_lists = {
(o, i): theta_np[o, i].reshape(-1).tolist()
for o in range(out_dim)
for i in range(in_dim)
}
x_py = x_np.tolist()
enc_py = encoded_x_np.tolist() if encoded_x_np is not None else None
for b in range(batch):
for o in range(out_dim):
for i in range(in_dim):
t = theta_lists[(o, i)]
enc_vals = enc_py[b][o][i] if enc_py is not None else None
if ansatz in ("pz_encoding", "pz"):
qc = _build_qiskit_pz_circuit(x_py[b][i], t, reps, enc_vals)
elif ansatz in ("rpz_encoding", "rpz"):
assert enc_vals is not None
qc = _build_qiskit_rpz_circuit(enc_vals, t, reps)
elif ansatz == "real":
qc = _build_qiskit_real_circuit(x_py[b][i], t, reps, enc_vals)
else:
raise NotImplementedError(
f"Ansatz '{ansatz}' not supported by qiskit_solver"
)
circuits.append(qc)
observables.append(pauli_z)
# Execute via the appropriate Estimator
max_pubs = config.get("max_pubs_per_job", 0)
mitigation = config.get("mitigation", {})
# Pre-build resources that don't change across ZNE/repeat calls
_rl = config.get("resilience_level")
_tw = config.get("twirling")
_max_execution_time = config.get("max_execution_time")
_job_tags = config.get("job_tags")
_job_event_callback = config.get("job_event_callback")
_pm = None
_rt_est = None
if (
backend is not None
and estimator is None
and not (parallel_qubits and parallel_qubits > 1)
):
_pm = generate_preset_pass_manager(
backend=backend,
optimization_level=optimization_level,
# Serial-path circuits are single-qubit: pin them to the
# best-ranked entry of the layout.
initial_layout=_initial_layout_for_circuit(initial_layout, 1),
)
_rt_est = Estimator(mode=backend)
_configure_estimator(
_rt_est,
shots,
_rl,
_tw,
max_execution_time=_max_execution_time,
job_tags=_job_tags,
)
execution_index = 0
def _run_qiskit(scale_factor=1):
nonlocal execution_index
job_context = {
"execution_index": execution_index,
"scale_factor": scale_factor,
}
execution_index += 1
run_circuits = (
[_fold_qiskit_circuit(qc, scale_factor) for qc in circuits]
if scale_factor > 1
else circuits
)
if parallel_qubits and parallel_qubits > 1:
return _qiskit_run_parallel(
run_circuits,
parallel_qubits,
estimator,
backend,
optimization_level,
shots,
initial_layout=initial_layout,
max_pubs_per_job=max_pubs,
resilience_level=_rl,
twirling=_tw,
max_execution_time=_max_execution_time,
job_tags=_job_tags,
job_event_callback=_job_event_callback,
job_context=job_context,
)
elif estimator is not None:
pubs = list(zip(run_circuits, observables))
context = _job_event_context(
job_context,
phase="direct",
pub_start=0,
pub_end=len(pubs),
pub_count=len(pubs),
logical_circuit_count=len(run_circuits),
)
job = _submit_job(
estimator,
pubs,
job_event_callback=_job_event_callback,
context=context,
)
result = _collect_job(
job,
job_event_callback=_job_event_callback,
context=context,
)
return [float(r.data.evs) for r in result]
elif _rt_est is not None:
isa_circuits = _pm.run(run_circuits)
isa_observables = [
obs.apply_layout(qc.layout)
for obs, qc in zip(observables, isa_circuits)
]
pubs = list(zip(isa_circuits, isa_observables))
context = _job_event_context(
job_context,
phase="direct",
pub_start=0,
pub_end=len(pubs),
pub_count=len(pubs),
logical_circuit_count=len(run_circuits),
)
job = _submit_job(
_rt_est,
pubs,
job_event_callback=_job_event_callback,
context=context,
)
result = _collect_job(
job,
job_event_callback=_job_event_callback,
context=context,
)
return [float(r.data.evs) for r in result]
else:
raise ValueError("No estimator or backend provided.")
if mitigation:
expvals = _apply_mitigation(_run_qiskit, mitigation)
else:
expvals = _run_qiskit(1)
output = torch.tensor(expvals, dtype=x.dtype, device=x.device)
return output.reshape(batch, out_dim, in_dim)
[docs]
def qiskit_solver(
x: torch.Tensor,
theta: torch.Tensor,
preacts_weight: torch.Tensor,
preacts_bias: torch.Tensor,
reps: int,
**kwargs,
) -> torch.Tensor:
"""
Execute QKAN circuits on IBM Quantum backends via Qiskit Runtime.
Drop-in replacement for torch_exact_solver. Circuits are built to match
the exact gate sequences of each ansatz, then executed on the specified
backend using Qiskit's Estimator primitive.
Supports training via the parameter-shift rule when gradients are needed.
Args
----
x : torch.Tensor
shape: (batch_size, in_dim)
theta : torch.Tensor
shape: (\\*group, reps+1, n_params) or (\\*group, reps, 1) for real
preacts_weight : torch.Tensor
shape: (\\*group, reps)
preacts_bias : torch.Tensor
shape: (\\*group, reps)
reps : int
ansatz : str
"pz_encoding", "pz", "rpz_encoding", "rpz", or "real"
preacts_trainable : bool
out_dim : int
backend : qiskit Backend
Qiskit backend instance (e.g., AerSimulator(), or from QiskitRuntimeService)
shots : int, optional
Number of shots per circuit. None for exact expectation (statevector).
optimization_level : int
Transpiler optimization level (0-3), default: 1
Returns
-------
torch.Tensor
shape: (batch_size, out_dim, in_dim)
"""
if not _QISKIT_AVAILABLE:
raise ImportError(
"Qiskit is required for qiskit_solver. "
"Install with: pip install qiskit qiskit-ibm-runtime"
)
ansatz = kwargs.get("ansatz", "pz_encoding")
preacts_trainable = kwargs.get("preacts_trainable", False)
out_dim = kwargs.get("out_dim", x.shape[1])
shots = kwargs.get("shots", None)
optimization_level = kwargs.get("optimization_level", 1)
parallel_qubits = kwargs.get("parallel_qubits", None)
backend = kwargs.get("backend", None)
estimator = kwargs.get("estimator", None)
# Resolve execution mode: estimator > backend > StatevectorEstimator > AerSimulator
if estimator is None and backend is None:
if _SV_ESTIMATOR_AVAILABLE:
estimator = StatevectorEstimator()
elif _AER_AVAILABLE:
backend = AerSimulator(method="statevector")
else:
raise ValueError(
"No backend or estimator specified. Install qiskit >= 1.0 "
"(for StatevectorEstimator), qiskit-aer, or qiskit-ibm-runtime."
)
# Auto-detect QPU size from backend if parallel_qubits="auto"
auto_parallel_qubits = parallel_qubits == "auto"
if auto_parallel_qubits and backend is not None:
parallel_qubits = backend.num_qubits
# Resolve initial_layout:
# None (default) -> let the transpiler choose
# "auto" -> top-N best-calibrated qubits via best_qubits()
# list[int] -> user-supplied physical qubit indices (used as-is)
initial_layout = kwargs.get("initial_layout", None)
if isinstance(initial_layout, tuple):
initial_layout = list(initial_layout)
elif initial_layout is not None and hasattr(initial_layout, "tolist"):
# Normalize numpy arrays / torch tensors so the string sentinel
# check and per-chunk truncation see a plain list.
initial_layout = list(initial_layout.tolist())
if isinstance(initial_layout, list) and not initial_layout:
# best_qubits returns [] when calibration is unavailable; qiskit
# treats an empty layout like None, so normalize it here too.
initial_layout = None
max_readout_error = kwargs.get("max_readout_error", None)
qubit_error_threshold = kwargs.get("qubit_error_threshold", None)
has_layout_threshold = (
max_readout_error is not None or qubit_error_threshold is not None
)
if isinstance(initial_layout, str):
if initial_layout != "auto":
raise ValueError(
"initial_layout must be None, 'auto', or a list of physical "
f"qubit indices; got {initial_layout!r}"
)
if backend is None:
raise ValueError(
"initial_layout='auto' requires a backend with properties() "
"to score qubit calibration."
)
n_layout = parallel_qubits if (parallel_qubits and parallel_qubits > 1) else 1
initial_layout = best_qubits(
backend,
n_layout,
max_readout_error=max_readout_error,
qubit_error_threshold=qubit_error_threshold,
# With parallel_qubits="auto" the packing width follows the
# calibration quality instead of failing the threshold check.
strict=not (auto_parallel_qubits and has_layout_threshold),
)
if not initial_layout:
# No calibration available — fall back to transpiler default.
warnings.warn(
"initial_layout='auto': backend exposes no calibration data; "
"falling back to the transpiler's default layout.",
stacklevel=2,
)
initial_layout = None
elif auto_parallel_qubits and has_layout_threshold:
parallel_qubits = len(initial_layout)
elif auto_parallel_qubits and isinstance(initial_layout, list):
parallel_qubits = len(initial_layout)
if initial_layout is not None and estimator is not None:
warnings.warn(
"initial_layout is ignored when execution goes through an "
"estimator rather than a backend: circuits are submitted "
"without qkan-side transpilation.",
stacklevel=2,
)
max_pubs_per_job = kwargs.get("max_pubs_per_job", 0)
config = {
"ansatz": ansatz,
"preacts_trainable": preacts_trainable,
"out_dim": out_dim,
"backend": backend,
"estimator": estimator,
"shots": shots,
"optimization_level": optimization_level,
"parallel_qubits": parallel_qubits,
"initial_layout": initial_layout,
"max_pubs_per_job": max_pubs_per_job,
"resilience_level": kwargs.get("resilience_level", None),
"twirling": kwargs.get("twirling", None),
"max_execution_time": kwargs.get("max_execution_time", None),
"job_tags": kwargs.get("job_tags", None),
"job_event_callback": kwargs.get("job_event_callback", None),
"mitigation": kwargs.get("mitigation", {}),
}
needs_grad = theta.requires_grad or x.requires_grad
if preacts_trainable:
needs_grad = (
needs_grad or preacts_weight.requires_grad or preacts_bias.requires_grad
)
if needs_grad:
return _QiskitParamShift.apply(
x, theta, preacts_weight, preacts_bias, reps, config
)
else:
return _qiskit_evaluate(x, theta, preacts_weight, preacts_bias, reps, config)
class QiskitSolver(QKANSolver):
"""Qiskit Runtime solver (registered as ``"qiskit"``)."""
name = "qiskit"
def __call__(
self,
x: torch.Tensor,
theta: torch.Tensor,
preacts_weight: torch.Tensor,
preacts_bias: torch.Tensor,
reps: int,
**kwargs,
) -> torch.Tensor:
return qiskit_solver(x, theta, preacts_weight, preacts_bias, reps, **kwargs)
register(QiskitSolver())