API Reference

This page contains the complete API reference for QKAN.

Core Module

class qkan.DARUAN(dim, reps, device='cpu', solver='exact', ansatz='pz_encoding', preact_trainable=False, postact_weight_trainable=False, postact_bias_trainable=False, seed=0, p_dim=4)

Bases: Module

forward(x)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class qkan.FourierKAN(layers_hidden, grid_size=5, scale_base=1.0, scale_fourier=1.0, base_activation=<class 'torch.nn.modules.activation.SiLU'>, grid_range=[-1, 1], device='cpu', seed=0, **kwargs)

Bases: Module

Fourier KAN (Kolmogorov-Arnold Network) model.

Uses Fourier basis (sin/cos) as the learnable activation functions on edges, serving as a baseline for comparison with B-spline KAN and QKAN.

forward(x, update_grid=False)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

regularization_loss(regularize_activation=1.0, regularize_entropy=1.0)

initialize_from_qkan(qkan, x0, sampling=100)

Initialize FourierKAN from a QKAN.

Parameters:

• qkan (QKAN)

• x0 (torch.Tensor (batch, in_dim))

• sampling (int)

class qkan.KAN(layers_hidden, grid_size=5, spline_order=3, scale_noise=0.1, scale_base=1.0, scale_spline=1.0, base_activation=<class 'torch.nn.modules.activation.SiLU'>, grid_eps=0.02, grid_range=[-1, 1], device='cpu', seed=0, **kwargs)

Bases: Module

KAN (Kolmogorov-Arnold Network) model. This is an efficient implementation of the KAN model, which is a neural network that uses B-spline as the learnable variational activation function.

It can be used for regression tasks and can be initialized from a QKAN model.

forward(x, update_grid=False)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

regularization_loss(regularize_activation=1.0, regularize_entropy=1.0)

initialize_from_qkan(qkan, x0, sampling=100)

Initialize KAN from a QKAN.

Parameters:

• qkan (QKAN)

• x0 (torch.Tensor (batch, in_dim))

• sampling (int)

class qkan.QKAN(width, reps=3, group=-1, is_map=False, is_batchnorm=False, hidden=0, device='cpu', solver='exact', qml_device='default.qubit', ansatz='pz_encoding', theta_size=None, norm_out=0, preact_trainable=False, preact_init=False, postact_weight_trainable=False, postact_bias_trainable=False, base_activation='silu', ba_trainable=False, fast_measure=True, save_act=False, c_dtype=torch.complex64, p_dtype=torch.float32, seed=None, solver_kwargs=None, p_dim=4, checkpoint_reps=False, theta_init='xavier', **kwargs)

Bases: Module

Quantum-inspired Kolmogorov Arnold Network (QKAN) Class

A quantum-inspired neural network that uses DatA Re-Uploading ActivatioN (DARUAN) as its learnable variation activation function.

References

Quantum Variational Activation Functions Empower Kolmogorov-Arnold Networks: https://arxiv.org/abs/2509.14026

width

List of width of each layer

Type:

list[int]

reps

Repetitions of quantum layers

Type:

int

group

Group of neurons

Type:

int

device

Device to use

Type:

Literal[“cpu”, “cuda”]

solver

Solver to use, currently supports “qml”, “exact”, “flash”, “cutn” or custom callable

Type:

Union[str, Callable]

qml_device

PennyLane device to use

Type:

str

layers

List of layers

Type:

QKANModuleList

is_map

Whether to use map layer

Type:

bool

is_batchnorm

Whether to use batch normalization

Type:

bool

reps

Repetitions of quantum layers

Type:

int

norm_out

Normalize output

Type:

int

postact_weight_trainable

Whether postact weights are trainable

Type:

bool

postact_bias_trainable

Whether postact bias are trainable

Type:

bool

preact_trainable

Whether preact weights are trainable

Type:

bool

base_activation

Base activation function

Type:

torch.nn.Module or lambda function

ba_trainable

Whether base activation weights are trainable

Type:

bool

fast_measure

Enable to use fast measurement in exact solver. Which would be quantum-inspired method. When False, the exact solver simulates the exact measurement process of quantum circuit.

Type:

bool

save_act

Whether to save activations

Type:

bool

seed

Random seed

Type:

int

to(*args, **kwargs)

Move the model to the specified device.

Parameters:

device (str | torch.device) – Device to move the model to, default: “cpu”

property param_size

forward(x)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

initialize_from_another_model(another_model)

Initialize from another model. Used for layer extension to refine the model.

Parameters:

another_model (QKAN) – Another model to initialize from

initialize_parameters()

Reinitialize parameters of all QKANLayer layers in-place.

xavier_init()

Apply Xavier normal initialization to all QKANLayer layers.

refine(new_reps)

Refine the model by layer extension, increasing the number of repetitions of quantum layers.

Parameters:

new_reps (int) – New number of repetitions of quantum layers

Return type:

QKAN

Returns:

QKAN

New QKAN model with increased repetitions

layer_extension(new_reps)

Refine the model by layer extension, increasing the number of repetitions of quantum layers.

Parameters:

new_reps (int) – New number of repetitions of quantum layers

Return type:

QKAN

Returns:

QKAN

New QKAN model with increased repetitions

get_reg(reg_metric, lamb_l1, lamb_entropy, lamb_coef, lamb_coefdiff)

Get regularization from the model.

Adapted from “pykan”.

Parameters:

• reg_metric (str) – Regularization metric. ‘edge_forward_dr_n’, ‘edge_forward_dr_u’, ‘edge_forward_sum’, ‘edge_backward’, ‘node_backward’

• lamb_l1 (float) – L1 Regularization parameter

• lamb_entropy (float) – Entropy Regularization parameter

• lamb_coef (float) – Coefficient Regularization parameter

• lamb_coefdiff (float) – Coefficient Smoothness Regularization parameter

Returns:

torch.Tensor

attribute(l=None, i=None, out_score=None, plot=True)

Get attribution scores

Adapted from “pykan”.

Parameters:

• l (None | int) – layer index

• i (None | int) – neuron index

• out_score (None | torch.Tensor) – specify output scores

• plot (bool) – when plot = True, display the bar show

Returns:

torch.Tensor

attribution scores

node_attribute()

Get node attribution scores.

Adapted from “pykan”.

train_(dataset, optimizer=None, closure=None, scheduler=None, steps=10, log=1, loss_fn=None, batch=-1, lamb=0.0, lamb_l1=1.0, lamb_entropy=2.0, lamb_coef=0.0, lamb_coefdiff=0.0, reg_metric='edge_forward_dr_n', verbose=True, clip_theta_grad_norm=None, lamb_theta_l1=0.0, prune_every=0, prune_threshold=0.001)

Train the model

Parameters:

• dataset (dict) – Dictionary containing train_input, train_label, test_input, test_label

• optimizer (torch.optim.Optimizer | None) – Optimizer to use, default: None

• closure (Callable | None) – Closure function for optimizer, default: None

• scheduler (torch.optim.lr_scheduler | None) – Scheduler to use, default: None

• steps (int) – Number of steps, default: 10

• log (int) – Logging frequency, default: 1

• loss_fn (torch.nn.Module | Callable |None) – Loss function to use, default: None

• batch (int) – batch size, if -1 then full., default: -1

• lamb (float) – L1 Regularization parameter. If 0, no regularization.

• lamb_l1 (float) – L1 Regularization parameter

• lamb_entropy (float) – Entropy Regularization parameter

• lamb_coef (float) – Coefficient Regularization parameter

• lamb_coefdiff (float) – Coefficient Smoothness Regularization parameter

• reg_metric (str) – Regularization metric. ‘edge_forward_dr_n’, ‘edge_forward_dr_u’, ‘edge_forward_sum’, ‘edge_backward’, ‘node_backward’

• verbose (bool) – Verbose mode, default: True

• clip_theta_grad_norm (float | None) – If not None, apply torch.nn.utils.clip_grad_norm_ to the concatenation of all QKANLayer.theta parameters (only) with this max_norm. Useful at high reps / bf16 where theta gradients occasionally explode near saddle points. Default None (no clipping).

• lamb_theta_l1 (float) – Weight for an additional sum(|theta|) penalty added to the training objective. Pushes low-contribution rotations toward zero so periodic mask pruning can zero them out. Default 0.0 (no penalty).

• prune_every (int) – If > 0, every prune_every steps, set mask[o,i]=0 wherever theta[o,i].abs().sum() < prune_threshold and flip _mask_is_identity=False on affected layers. Layers with grouped (shared) theta are skipped — no per-edge magnitude exists. Default 0 (no auto-pruning).

• prune_threshold (float) – Threshold for the per-edge theta-magnitude check used by prune_every. Default 1e-3.

Returns:

dict

Dictionary containing train_loss and test_loss

plot(x0=None, sampling=1000, from_acts=False, scale=0.5, beta=3, metric='forward_n', mask=False, in_vars=None, out_vars=None, title=None)

Plot the model.

Adapted from “pykan”.

Parameters:

• x0 (torch.Tensor | None) – Input tensor to plot, if None, plot from saved activations

• sampling (int) – Sampling frequency

• from_acts (bool) – Plot from saved activations

• scale (float) – Scale of the plot

• beta (float) – Beta value

• metric (str) – Metric to use. ‘forward_n’, ‘forward_u’, ‘backward’

• in_vars (list[int] | None) – Input variables to plot

• out_vars (list[int] | None) – Output variables to plot

• title (str | None) – Title of the plot

prune_node(threshold=0.01, mode='auto', active_neurons_id=None)

Pruning nodes.

Adapted from “pykan”.

Parameters:

• threshold (float) – if the attribution score of a neuron is below the threshold, it is considered dead and will be removed

• mode (str) – “auto” or “manual”. with “auto”, nodes are automatically pruned using threshold. With “manual”, active_neurons_id should be passed in.

Returns:

QKAN

pruned network

prune_edge(threshold=0.03)

Pruning edges.

Adapted from “pykan”.

Parameters:

threshold (float) – float if the attribution score of an edge is below the threshold, it is considered dead and will be set to zero.

prune(node_th=0.01, edge_th=0.03)

Prune (both nodes and edges).

Adapted from “pykan”.

Parameters:

• node_th (float) – if the attribution score of a node is below node_th, it is considered dead and will be set to zero.

• edge_th (float) – if the attribution score of an edge is below node_th, it is considered dead and will be set to zero.

Returns:

QKAN

pruned network

prune_input(threshold=0.01, active_inputs=None)

Prune inputs.

Adapted from “pykan”.

Parameters:

• threshold (float) – if the attribution score of the input feature is below threshold, it is considered irrelevant.

• active_inputs (list | None) – if a list is passed, the manual mode will disregard attribution score and prune as instructed.

Returns:

QKAN

pruned network

remove_edge(layer_idx, in_idx, out_idx)

Remove activtion phi(layer_idx, in_idx, out_idx) (set its mask to zero)

Parameters:

• layer_idx (int) – Layer index

• in_idx (int) – Input node index

• out_idx (int) – Output node index

remove_node(layer_idx, in_idx, mode='all')

remove neuron (layer_idx, in_idx) (set the masks of all incoming and outgoing activation functions to zero)

Parameters:

• layer_idx (int) – Layer index

• in_idx (int) – Input node index

• mode (str) – Mode to remove. “all” or “up” or “down”, default: “all”

static clear_ckpts(folder='./model_ckpt')

Clear all checkpoints.

Parameters:

folder (str) – Folder containing checkpoints, default: “./model_ckpt”

save_ckpt(name, folder='./model_ckpt')

Save the current model as checkpoint.

Parameters:

• name (str) – Name of the checkpoint

• folder (str) – Folder to save the checkpoint, default: “./model_ckpt”

load_ckpt(name, folder='./model_ckpt')

Load a checkpoint to the current model.

Parameters:

• name (str) – Name of the checkpoint

• folder (str) – Folder containing the checkpoint, default: “./model_ckpt”

class qkan.AdaBelief(params, lr=0.01, betas=(0.9, 0.999), eps=1e-16, weight_decay=0.0)

Bases: Optimizer

AdaBelief — drop-in Adam replacement with variance-of-gradient s.

Parameters:

• params (Iterable[Any]) – iterable of parameters.

• lr (float) – learning rate. Default 1e-2 — for QKAN this is materially different from the standard Adam 1e-3; sweep on your task.

• betas (tuple[float, float]) – (β₁, β₂). Defaults match the paper.

• eps (float) – numerical floor added to √ŝ in the denominator.

• weight_decay (float) – decoupled (AdamW-style) weight decay.

step(closure=None)

Performs a single optimization step (parameter update).

Parameters:

closure (Callable) – A closure that reevaluates the model and returns the loss. Optional for most optimizers.

Return type:

Optional[float]

Note

Unless otherwise specified, this function should not modify the .grad field of the parameters.

class qkan.QKANAdamMini(params, lr=0.001, betas=(0.9, 0.999), eps=1e-08, weight_decay=0.0)

Bases: Optimizer

Adam-mini optimizer with QKAN-aware block partitioning.

Parameters:

• params (iterable) – Iterable of parameters to optimize, or dicts defining parameter groups. Parameters passed as (name, parameter) tuples (via model.named_parameters()) get block layouts inferred from the name; otherwise all parameters get per-tensor blocking.

• lr (float) – Learning rate.

• betas ((float, float)) – Coefficients for first and second moment running averages.

• eps (float) – Numerical stability term added to sqrt(v).

• weight_decay (float) – Decoupled weight decay coefficient (AdamW-style).

Notes

For best results, build with QKANAdamMini(model.named_parameters(), ...) — passing names lets the optimizer pick per-edge blocks for theta / preacts. Without names, it falls back to per-tensor blocking (still correct, just with less memory savings).

step(closure=None)

Performs a single optimization step (parameter update).

Parameters:

closure (Callable) – A closure that reevaluates the model and returns the loss. Optional for most optimizers.

Return type:

Optional[float]

Note

Unless otherwise specified, this function should not modify the .grad field of the parameters.

class qkan.QKANBeliefMini(params, lr=0.01, betas=(0.9, 0.999), eps=1e-16, weight_decay=0.0, state_dtype=None)

Bases: Optimizer

AdaBelief with Adam-mini block partitioning for s.

The first moment m stays per-parameter (same as Adam-mini / AdaBelief); the variance s collapses to one scalar per Adam-mini block, partitioned by the same rule as QKANAdamMini:

• theta natural (O, I, R+1, K) → block per (o, i, r)

• preacts_* (O, I, R) → block per (o, i)

• (O, I) params → one block per tensor

• non-QKAN Linear weights → per output row

• other / LayerNorm / 1-D → one block per tensor

Optimizer state ~30% smaller than full AdaBelief; convergence sits between Adam and AdaBelief.

Pass model.named_parameters() (with names) to get the QKAN detection; bare model.parameters() falls back to per-tensor.

Parameters:

• params (Iterable[Any]) – iterable of parameters (or (name, param) tuples).

• lr (float) – learning rate.

• betas (tuple[float, float]) – (β₁, β₂).

• eps (float) – numerical floor on the denominator.

• weight_decay (float) – decoupled (AdamW-style) weight decay.

• state_dtype (Optional[dtype]) – dtype for m and the block-reduced s. None (default) inherits the param dtype. Pass torch.bfloat16 to halve state memory; compute stays native (torch’s add/mul handle bf16 correctly enough for these tiny accumulators). Note: when the params themselves are bf16 and state_dtype=None, s will accumulate squared residuals in bf16 and may underflow on long runs — pass state_dtype=torch.float32 explicitly to be safe.

step(closure=None)

Performs a single optimization step (parameter update).

Parameters:

closure (Callable) – A closure that reevaluates the model and returns the loss. Optional for most optimizers.

Return type:

Optional[float]

Note

Unless otherwise specified, this function should not modify the .grad field of the parameters.

class qkan.QKANMuon(params, lr=0.02, adamw_lr=0.003, momentum=0.95, nesterov=True, ns_steps=5, weight_decay=0.0, match_adamw_rms=True, adamw_betas=(0.9, 0.999), adamw_eps=1e-08, muon_filter=None, rank=None, world_size=None, process_group=None)

Bases: Optimizer

Muon for the 2-D matrices of a QKAN-bearing transformer; AdamW elsewhere.

Parameters:

• params (iterable) – Parameters or (name, parameter) tuples from named_parameters(). Names drive the Muon/AdamW routing.

• lr (float) – Learning rate for the Muon (matrix) path.

• adamw_lr (float) – Learning rate for the AdamW (fallback) path.

• momentum (float) – Muon momentum.

• nesterov (bool) – Nesterov-style lookahead on the Muon momentum.

• ns_steps (int) – Newton-Schulz iterations.

• weight_decay (float) – Decoupled weight decay (applied on both paths).

• match_adamw_rms (bool) – If True, scale the orthogonalized update by 0.2*sqrt(max(out,in)) so its per-element RMS matches AdamW’s (Moonlight). If False, use Keller’s max(1, out/in)**0.5 aspect-ratio scale (vanilla Muon).

• adamw_betas (tuple[float, float]) – AdamW hyperparameters (defaults match torch.optim.AdamW).

• adamw_eps (float) – AdamW hyperparameters (defaults match torch.optim.AdamW).

• muon_filter (optional callable (name, shape) -> bool) – Overrides the default routing (used for ablations and tests).

• rank (int, optional) – Distributed rank and world size. If omitted, auto-detected from torch.distributed when initialized, else single-process (0, 1).

• world_size (int, optional) – Distributed rank and world size. If omitted, auto-detected from torch.distributed when initialized, else single-process (0, 1).

• process_group (optional) – Process group for the collective; defaults to the world group. Its ranks are assumed to be 0..world_size-1 (the standard DDP layout).

Notes

Distributed (world_size > 1) follows the nanochat / Keller-Jordan sharding pattern: gradients are assumed already all-reduced across ranks (e.g. via DDP), then the Muon matrices are round-robined across ranks so each rank runs the expensive Newton-Schulz on only 1/world_size of them; each updated matrix is then broadcast from its owner. The broadcasts are synchronous (not overlapped with compute). Momentum buffers live only on the owning rank (sharded optimizer state), while the AdamW-fallback params stay replicated (every rank applies the identical update on the all-reduced grad). world_size == 1 issues no collective and is byte-identical to the single-process path.

Because Muon momentum is rank-sharded, state_dict() holds only this rank’s shard: checkpoint/restore requires every rank to save and load its own state_dict, and resuming requires the same world_size.

static is_muon_param(name, shape)

Default routing: True -> Muon path, False -> AdamW path.

Decided purely from the parameter name and rank, which are invariant to QKAN’s p_dim storage — so (unlike the Adam-mini-family siblings) this needs no _qkan_natural_shape: any 2-D QKAN param (base_weight, postact, or a p_dim=2-collapsed theta) is routed to AdamW by its name marker regardless of its stored shape.

Return type:

bool

describe_layout()

Return (name, shape, "muon"|"adamw") per parameter.

Return type:

list[tuple[str, tuple[int, ...], str]]

step(closure=None)

Performs a single optimization step (parameter update).

Parameters:

closure (Callable) – A closure that reevaluates the model and returns the loss. Optional for most optimizers.

Return type:

Optional[float]

Note

Unless otherwise specified, this function should not modify the .grad field of the parameters.

class qkan.QKANLayer(in_dim, out_dim, reps=3, group=-1, device='cpu', solver='exact', qml_device='default.qubit', ansatz='pz_encoding', theta_size=None, preact_trainable=False, preact_init=False, postact_weight_trainable=False, postact_bias_trainable=False, base_activation='silu', ba_trainable=True, is_batchnorm=False, fast_measure=True, c_dtype=torch.complex64, p_dtype=torch.float32, seed=None, solver_kwargs=None, p_dim=4, checkpoint_reps=False, theta_init='xavier')

Bases: Module

QKANLayer Class

in_dim

Input dimension

Type:

int

out_dim

Output dimension

Type:

int

reps

Repetitions of quantum layers

Type:

int

group

Group of neurons

Type:

int

device

Device to use

solver

Solver to use, currently supports “qml”, “exact”, “flash”, “cutn” or custom callable

Type:

Union[str, Callable]

ansatz

Ansatz to use, “pz_encoding”, “px_encoding”, “rpz_encoding” or custom

Type:

Union[str, Callable]

qml_device

PennyLane device to use

Type:

str

theta

Learnable parameter of quantum circuit

Type:

nn.Parameter

base_weight

Learnable parameter of base activation

Type:

nn.Parameter

preact_trainable

Whether preact weights are trainable

Type:

bool

preacts_weight

Learnable parameter of preact weights

Type:

nn.Parameter

preacts_bias

Learnable parameter of preact bias

Type:

nn.Parameter

postact_weight_trainable

Whether postact weights are trainable

Type:

bool

postact_weights

Learnable parameter of postact weights

Type:

nn.Parameter

postact_bias_trainable

Whether postact bias are trainable

Type:

bool

postact_bias

Learnable parameter of postact bias

Type:

nn.Parameter

mask

Mask for pruning

Type:

nn.Parameter

is_batchnorm

Whether to use batch normalization

Type:

bool

fast_measure

Enable to use fast measurement in exact solver. Which would be quantum-inspired method. When False, the exact solver simulates the exact measurement process of quantum circuit.

Type:

bool

c_dtype

Compute dtype for quantum simulation. Supported values:

• torch.complex64 / torch.float32: full-precision f32 (default)

• torch.bfloat16: mixed-precision bf16 I/O, f32 compute, bf16 state checkpoints

• torch.float8_e4m3fn: bf16 I/O, f32 compute, fp8 prescaled state checkpoints

Type:

torch.dtype

p_dtype

Parameter dtype (torch.float32 or torch.bfloat16). Use torch.bfloat16 with bf16/fp8 c_dtype for full mixed-precision pipeline.

Type:

torch.dtype

_x0

Leave for ResQKANLayer

Type:

Optional[torch.Tensor]

init_parameters()

Create all learnable parameters.

Called once from __init__ to allocate nn.Parameter objects. Reads configuration from self.* attributes. If self.seed is set, the RNG is seeded for reproducibility.

Calls xavier_init() at the end to apply Xavier normal initialization to theta (and preacts when preact_init is set).

xavier_init()

Initialise theta and (optionally) preacts.

Theta init follows self.theta_init:

• "xavier" (default): Xavier-normal init (fan-in/fan-out aware).

• "small_gaussian": theta ~ N(0, sigma^2) with sigma = 1/sqrt(reps). Mitigates the barren-plateau regime in deep parametric quantum circuits (Zhang+ 2022 arXiv:2203.09376). Opt-in for reps >= 8; underperforms xavier at small reps.

Preacts (when preact_init is True) keep Xavier — they are amplitude-scaled so fan-in/out are meaningful.

to(*args, **kwargs)

Move the layer to the specified device.

Parameters:

device (str | torch.device) – Device to move the layer to, default: “cpu”

train(mode=True)

Set the module in training mode.

This has any effect only on certain modules. See documentations of particular modules for details of their behaviors in training/evaluation mode, if they are affected, e.g. Dropout, BatchNorm, etc.

Parameters:

mode (bool) – whether to set training mode (True) or evaluation mode (False). Default: True.

Returns:

self

Return type:

Module

property param_size

property x0

forward(x)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

set_fused_epilogue(enabled=True)

Toggle the fused Triton epilogue path (forward + backward).

Off by default. When on, _forward_train / _forward_eval collapse the (postacts + pb)·eff_pw sum and the base linear term into a single kernel launch; the backward fuses grad_postacts / grad_eff_pw / grad_pb into one Triton kernel and keeps the matmul-shaped gradients on cuBLAS. CPU tensors and non-CUDA paths transparently fall back to the eager chain.

Return type:

None

reset_parameters()

Reset all learnable parameters to default values in-place.

Note: The thetas are set to zero to do layer extension. If you wish to re-init the parameters, please use init_parameters instead.

forward_no_sum(x)

get_subset(in_id, out_id)

Get a smaller QKANLayer from a larger QKANLayer (used for pruning).

Parameters:

• in_id (list) – id of selected input neurons

• out_id (list) – id of selected output neurons

Returns:

QKANLayer

New QKANLayer with selected neurons

class qkan.QKANSpectralMini(params, lr=0.001, betas=(0.9, 0.999), eps=0.001, weight_decay=0.0)

Bases: Optimizer

Eigenvalue-aware Adam-mini using a rank-1 GN preconditioner.

Parameters:

• params (iterable) – Iterable of parameters or (name, parameter) tuples from model.named_parameters(). Names enable per-block layout.

• lr (float) – Learning rate.

• betas ((float, float)) – beta1 for momentum, beta2 for the GN-direction EMA.

• eps (float) – Damping lambda for the Sherman-Morrison inverse. Acts like the diagonal floor of the curvature estimate. Default 1e-3 empirically stable on QKAN.

• weight_decay (float) – AdamW-style decoupled weight decay.

describe_layout()

Return (name, shape, block_ndim, num_blocks) per parameter.

Useful for sanity-checking the block partition at construction.

Return type:

list[tuple[str, tuple[int, ...], int, int]]

step(closure=None)

Performs a single optimization step (parameter update).

Parameters:

closure (Callable) – A closure that reevaluates the model and returns the loss. Optional for most optimizers.

Return type:

Optional[float]

Note

Unless otherwise specified, this function should not modify the .grad field of the parameters.

class qkan.StateVector(batch_size, out_dim, in_dim, device='cpu', dtype=torch.complex64)

Bases: object

1-qubit state vector.

StateVector.state: torch.Tensor, shape: (batch_size, out_dim, in_dim, 2)

state: Tensor

measure_z(fast_measure=True)

Measure the state vector in the Z basis.

Return type:

Tensor

:paramif False, return |α|^2 - |β|^2.

Which is quantum-inspired method and faster when it is True.

:type : fast_measure: bool, default: True. If True, for state |ψ⟩ = α|0⟩ + β|1⟩, return |α| - |β|;

return: torch.Tensor, shape: (batch_size, out_dim, in_dim)

measure_x(fast_measure=True)

Measure the state vector in the X basis.

Return type:

Tensor

:paramif False, return |α|^2 - |β|^2.

Which is quantum-inspired method and faster when it is True.

:type : fast_measure: bool, default: True. If True, for state |ψ⟩ = α|0⟩ + β|1⟩, return |α| - |β|;

return: torch.Tensor, shape: (batch_size, out_dim, in_dim)

measure_y(fast_measure=True)

Measure the state vector in the Y basis.

Return type:

Tensor

:paramif False, return |α|^2 - |β|^2.

Which is quantum-inspired method and faster when it is True.

:type : fast_measure: bool, default: True. If True, for state |ψ⟩ = α|0⟩ + β|1⟩, return |α| - |β|;

return: torch.Tensor, shape: (batch_size, out_dim, in_dim)

s(is_dagger=False)

Apply Phase gate (or S gate) to the state vector.

:param : :type : is_dagger: bool, default: False

h(is_dagger=False)

Apply Hadamard gate to the state vector.

:param : :type : is_dagger: bool, default: False

x()

Apply Pauli-X gate to the state vector.

z()

Apply Pauli-Z gate to the state vector.

rx(theta, is_dagger=False)

Apply Rotation-X gate to the state vector.

:param : :type : theta: torch.Tensor, shape: (out_dim, in_dim) :param : :type : is_dagger: bool, default: False

ry(theta, is_dagger=False)

Apply Rotation-Y gate to the state vector.

:param : :type : theta: torch.Tensor, shape: (out_dim, in_dim) :param : :type : is_dagger: bool, default: False

rz(theta, is_dagger=False)

Apply Rotation-Z gate to the state vector.

:param : :type : theta: torch.Tensor, shape: (out_dim, in_dim) :param : :type : is_dagger: bool, default: False

class qkan.TorchGates

Bases: object

static identity_gate(shape)

shape: (out_dim, in_dim)

return: torch.Tensor, shape: (2, 2, out_dim, in_dim)

Return type:

Tensor

static i_gate(shape)

shape: (out_dim, in_dim)

return: torch.Tensor, shape: (2, 2, out_dim, in_dim)

Return type:

Tensor

static rx_gate(theta, dtype=torch.complex64)

theta: torch.Tensor, shape: (out_dim, in_dim)

return: torch.Tensor, shape: (2, 2, out_dim, in_dim)

Return type:

Tensor

static ry_gate(theta, dtype=torch.complex64)

theta: torch.Tensor, shape: (out_dim, in_dim)

return: torch.Tensor, shape: (2, 2, out_dim, in_dim)

Return type:

Tensor

static rz_gate(theta, dtype=torch.complex64)

theta: torch.Tensor, shape: (out_dim, in_dim)

return: torch.Tensor, shape: (2, 2, out_dim, in_dim)

Return type:

Tensor

static h_gate(shape, device, dtype=torch.complex64)

shape: (out_dim, in_dim)

return: torch.Tensor, shape: (2, 2, out_dim, in_dim)

Return type:

Tensor

static s_gate(shape)

shape: (out_dim, in_dim)

return: torch.Tensor, shape: (2, 2, out_dim, in_dim)

Return type:

Tensor

static acrx_gate(theta, dtype=torch.complex64)

Complex extension of RX(acos(theta)) gate. Note: Physically unrealizable.

theta: torch.Tensor, shape: (out_dim, in_dim)

return: torch.Tensor, shape: (2, 2, out_dim, in_dim)

Return type:

Tensor

static tensor_product(gate, another_gate, dtype=None)

Compute tensor product of two gates.

:param : :type : gate: torch.Tensor, shape: (2, 2, out_dim, in_dim) :param : :type : another_gate: torch.Tensor, shape: (2, 2, out_dim, in_dim) :param : Both gates should have the same dtype. :type : dtype: torch dtype, optional. If None, uses the dtype of the input gate.

return: torch.Tensor, shape: (4, 4, out_dim, in_dim)

static cx_gate(shape, control, device, dtype=torch.complex64)

2-qubits CX (CNOT) gate.

shape: (out_dim, in_dim) control: int

return: torch.Tensor, shape: (4, 4, out_dim, in_dim)

Return type:

Tensor

static cz_gate(shape, device, dtype=torch.complex64)

2-qubits CZ gate.

shape: (out_dim, in_dim) control: int

return: torch.Tensor, shape: (4, 4, out_dim, in_dim)

Return type:

Tensor

class qkan.CompiledInference(module, max_shapes=8, warmup=3)

Bases: Module

Transparent wrapper that captures CUDA graphs lazily per input shape.

Behaves exactly like the wrapped module in training (self.training is True) or when grad is enabled or the input is not a single CUDA tensor — so you can wrap a model once and keep using it normally:

model = CompiledInference(model)
model.train(); model(x).sum().backward()   # eager, grad flows
model.eval()
with torch.no_grad():
    model(x)                               # captures + replays

On each eval/no-grad forward, the (shape, dtype, device) of the input is used as the cache key. A miss triggers a capture (up to max_shapes); a hit replays the captured graph. When the cache is full the fallback is eager execution — graphs are not evicted.

Parameters:

• module (Module) – the module to wrap. forward must take one tensor arg.

• max_shapes (int) – maximum number of distinct input shapes to cache.

• warmup (int) – warmup forward passes before each capture.

train(mode=True)

Set the module in training mode.

This has any effect only on certain modules. See documentations of particular modules for details of their behaviors in training/evaluation mode, if they are affected, e.g. Dropout, BatchNorm, etc.

Parameters:

mode (bool) – whether to set training mode (True) or evaluation mode (False). Default: True.

Returns:

self

Return type:

Module

clear_cache()

Drop all captured graphs. Call after editing parameters in-place.

Return type:

None

forward(x, *extra, **kwargs)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

Return type:

Tensor

class qkan.DeviceProfile(num_qubits, edges=(), readout_error=<factory>, gate_error_1q=<factory>, gate_error_2q=<factory>, t2_us=None, operational=None)

Bases: object

Provider-neutral snapshot of device calibration.

Physical qubit labels run 0..num_qubits-1; devices with non-contiguous or 1-indexed labels simply have holes in the maps, and unreported labels are marked non-operational by the adapters. Coupling edges and the gate_error_2q map are keyed by sorted (low, high) pairs; t2_us is in microseconds. All error values are probabilities in [0, 1]; NaN entries are kept and treated as missing data at scoring time, and missing dict keys mean the device reported nothing for that qubit or edge (common for faulty qubits). Qubits mapped to False in operational are never selected.

num_qubits: int

edges: tuple[tuple[int, int], ...] = ()

readout_error: Mapping[int, float]

gate_error_1q: Mapping[int, float]

gate_error_2q: Mapping[tuple[int, int], float]

t2_us: Mapping[int, float] | None = None

operational: Mapping[int, bool] | None = None

has_calibration()

True when the profile carries any per-qubit calibration data.

Return type:

bool

classmethod from_qiskit(backend, *, refresh=False)

Build a profile from a qiskit backend.

Prefers the BackendV2 Target (structured per-instruction errors); falls back to the legacy properties() API. With refresh=True the backend is asked to re-fetch calibration first — via backend.refresh() on BackendV2 (qiskit-ibm-runtime) or properties(refresh=True) on the legacy path — since both APIs cache their first snapshot.

Return type:

DeviceProfile

classmethod from_braket(device)

Build a profile from an AWS Braket device (experimental).

device is a braket.aws.AwsDevice (or any object exposing the same properties model). Reads the standardized gate-model QPU calibration schema: per-qubit T1/T2, readout fidelity (or IQM’s readout-error entries), and 1Q (simultaneous) randomized-benchmarking fidelity, plus per-pair 2Q gate fidelities. Fidelities are converted to errors as 1 - fidelity; entries whose type name starts with READOUT_ERROR are already errors and are averaged as-is.

Return type:

DeviceProfile

class qkan.PackedCircuit(tiles, circuit=None, isa=None, parameters=(), kernel=None, gates=None, block=None)

Bases: object

A packed job: copies circuit copies pinned to disjoint tiles.

One class serves both stacks — pack_circuit() fills the half that applies, and the placement accessors are shared.

qiskit (estimator-driven): circuit is the merged logical circuit (copy t on logical qubits t*block_qubits ..) and isa its transpilation pinned to layout; map observables per tile with observable() and bind parameter batches with parameter_values() (each copy carries its own renamed parameters, parameters[t] in the source circuit.parameters order) or rebind().

CUDA-Q (counts-driven): kernel applies the block at each tile’s physical indices and (in automatic mode) measures the full register; sample it and read per-tile results with expectation(), X/Y observables through basis_kernel() (hardware-safe) or observable() (simulator observe), and bind argument batches with rebind(). gates holds one extracted gate list per tile and block the source kernel.

tiles: tuple[tuple[int, ...], ...]

circuit: Any = None

isa: Any = None

parameters: tuple[tuple[str, ...], ...] = ()

kernel: Any = None

gates: tuple | None = None

block: Any = None

property layout: tuple[int, ...]

Tile-major flattening of tiles (the pinned initial layout).

property width: int

Register width of the packed CUDA-Q kernel (max index + 1).

property block_qubits: int

Qubits per circuit copy.

property copies: int

Number of packed circuit copies.

physical_qubits(tile)

Physical qubit indices hosting circuit copy tile.

Return type:

list[int]

parameter_values(batch)

Bind one batch of block parameters — copy t gets batch[t].

batch[t][j] binds the j-th parameter (in the source circuit’s circuit.parameters order) of copy t. Returns the name-to-value mapping for an EstimatorV2 PUB alongside isa, so a variational loop packs once and re-binds every step.

Return type:

dict[str, float]

observable(obs, tile=None)

Map a block-level observable onto the packed placement.

Returns the mapped operator for one tile, the list over all copies (None), or their average ("mean") — when every tile runs the same circuit, the mean pools the tiles into one estimate at copies-fold shots.

qiskit packs take a Pauli string (little-endian) or SparsePauliOp over the block’s block_qubits qubits and return the ISA-mapped operator, ready for an EstimatorV2 PUB alongside isa. CUDA-Q packs take a Pauli string in index order (pauli[i] acts on block qubit i) and return a cudaq.spin operator at the tile’s physical indices, for cudaq.observe(packed.observe_kernel(), ...). That is the simulator route — hardware targets that compact idle qubits reject sparse-index observables; use basis_kernel() plus expectation() there.

rebind(batch)

Re-pack the same tiles with a new batch of values.

The per-step call of a variational loop — tile selection is reused. CUDA-Q packs re-extract the block per tile (batch holds one runtime-argument set per tile, as block_args_batch in pack_circuit(); milliseconds per step). qiskit packs return a copy with the parameters bound (batch rows as in parameter_values(); prefer parameter_values in estimator PUBs, which keeps the circuit parametric).

Return type:

PackedCircuit

basis_kernel(bases)

Packed kernel with per-qubit basis rotations before measurement.

bases has one letter per block qubit, position i acting on block qubit i: X and Y append that basis change (H, or RZ(-pi/2) then H) on the corresponding physical qubit of every tile; Z and I leave it untouched. Sample the returned kernel and read Pauli expectations with expectation() — the hardware-safe observable path (sparse-index observe fails on targets that compact idle qubits).

expectation(result, pauli, tile=None)

Pauli expectation values from a sample of the matching basis.

pauli[i] acts on block qubit i; result must come from sampling basis_kernel() for that string (the plain kernel suffices when it only contains Z and I). The value is the Z-parity of the non-identity positions — for one tile, the list over all copies (None), or their average ("mean"): when every tile runs the same circuit, the mean pools the tiles into one estimate at copies-fold shots.

observe_kernel()

Measurement-free packed kernel for cudaq.observe.

observe rejects kernels with measurements, so the sampled kernel (which measures the full register) cannot be used with observable(); this rebuild omits the measurement. Simulator route only — without the full-register measurement, hardware pipelines may compact idle qubits.

class qkan.LBFGSFinisher(early, params, total_steps, pct_early=0.7, lbfgs_kwargs=None)

Bases: object

Composite optimizer: early for pct_early * total_steps, then LBFGS.

Use step(closure) exactly like a regular torch.optim.Optimizer. The closure must zero grads, run forward+backward, and return loss:

def closure():
    opt.zero_grad()
    loss = loss_fn(model(x), y)
    loss.backward()
    return loss

for _ in range(total_steps):
    loss = opt.step(closure)

Parameters:

• early (torch.optim.Optimizer) – Any standard optimizer (Adam, QKANAdamMini, …). Will be used for the first pct_early of steps.

• params (iterable of nn.Parameter) – The same params passed to early. Needed to construct the L-BFGS optimizer at swap time.

• total_steps (int) – Total budget. After int(pct_early * total_steps) calls to step, the wrapper switches to L-BFGS for the remainder.

• pct_early (float) – Fraction of total_steps to run with early. Default 0.7 matches the pykan recipe.

• lbfgs_kwargs (dict, optional) – Override defaults for torch.optim.LBFGS. Sensible defaults for KAN-style fits: lr=1.0, max_iter=20, history_size=100, tolerance_grad=1e-7, line_search_fn='strong_wolfe'.

property using_lbfgs: bool

True if the next step will use L-BFGS.

property current: Optimizer

The optimizer that will handle the next step call.

zero_grad(set_to_none=True)

Zero gradients on the active optimizer.

Most users won’t call this directly — the closure does it. Provided for parity with torch.optim.Optimizer.

Return type:

None

step(closure)

Run one optimization step. closure is required.

L-BFGS requires a closure that zeros grads, runs forward+backward, and returns the loss tensor. The early optimizer is also driven via the same closure so user code stays uniform across the swap.

Return type:

Tensor

state_dict()

Return type:

dict[str, Any]

load_state_dict(state)

Return type:

None

class qkan.TritonAdaBelief(params, lr=0.01, betas=(0.9, 0.999), eps=1e-16, weight_decay=0.0, state_dtype=None)

Bases: Optimizer

Single-kernel AdaBelief step via Triton — same algorithm as AdaBelief.

Parameters:

• params (Iterable[Any]) – iterable of parameters.

• lr (float) – learning rate.

• betas (tuple[float, float]) – (β₁, β₂).

• eps (float) – numerical floor on the denominator.

• weight_decay (float) – decoupled (AdamW-style) weight decay.

• state_dtype (Optional[dtype]) – dtype for m and s. None (default) inherits the param dtype. Pass torch.bfloat16 to halve state memory; the kernel always computes in fp32 via implicit upcasts on load. Note: when the params themselves are bf16 and state_dtype=None, s will accumulate squared residuals in bf16 and may underflow on long runs — pass state_dtype=torch.float32 explicitly to be safe.

step(closure=None)

Performs a single optimization step (parameter update).

Parameters:

closure (Callable) – A closure that reevaluates the model and returns the loss. Optional for most optimizers.

Return type:

Optional[float]

Note

Unless otherwise specified, this function should not modify the .grad field of the parameters.

qkan.adam_then_lbfgs(model, total_steps, lr_adam=0.01, pct_adam=0.7, use_adam_mini=True, lbfgs_kwargs=None)

Convenience factory: build an Adam/AdamMini -> LBFGS schedule.

Parameters:

• model (nn.Module) – The model whose parameters will be optimized.

• total_steps (int) – Total budget. Adam runs for int(pct_adam * total_steps) steps, L-BFGS for the rest.

• lr_adam (float) – Learning rate for the Adam (or QKANAdamMini) phase.

• pct_adam (float) – Fraction of total_steps to spend in Adam. Default 0.7.

• use_adam_mini (bool) – If True (default), use QKANAdamMini as the early optimizer — block-aware, ~26% less optimizer state. If False, fall back to plain torch.optim.Adam.

• lbfgs_kwargs (dict, optional) – Override L-BFGS defaults.

Returns:

Ready to drive with opt.step(closure).

Return type:

LBFGSFinisher

qkan.compile_inference(module, max_shapes=8, warmup=3)

Shortcut: compile_inference(m) == CompiledInference(m).

Drop-in replacement that uses CUDA graphs for inference and falls back to eager execution for training or gradient-tracking forwards. See CompiledInference for details.

Return type:

CompiledInference

qkan.create_dataset(f, n_var=2, f_mode='col', ranges=[-1, 1], train_num=1000, test_num=1000, normalize_input=False, normalize_label=False, device='cpu', seed=0)

Create dataset

Parameters:

• f – function the symbolic formula used to create the synthetic dataset

• ranges – list or np.array; shape (2,) or (n_var, 2) the range of input variables. Default: [-1,1].

• train_num – int the number of training samples. Default: 1000.

• test_num – int the number of test samples. Default: 1000.

• normalize_input – bool If True, apply normalization to inputs. Default: False.

• normalize_label – bool If True, apply normalization to labels. Default: False.

• device – str device. Default: ‘cpu’.

• seed – int random seed. Default: 0.

Returns:

dict

Train/test inputs/labels are dataset[‘train_input’], dataset[‘train_label’], dataset[‘test_input’], dataset[‘test_label’]

Return type:

dataset

qkan.get_feynman_dataset(name)

Get Feynman dataset from the given name.

Parameters:

name (str | int) – The name of the dataset.

Returns:

The symbol, expression, function, and ranges of the dataset.

Return type:

tuple

qkan.tile_disjoint(profile, interaction, k='max', *, buffer_hops=0, max_readout_error=None, qubit_error_threshold=None, edge_error_threshold=None, tile_score_threshold=None, strict=True, call_limit=50000, max_trials=500, n_logical=None)

Tile the chip with disjoint calibration-aware copies of one circuit.

Greedy peel: repeatedly place the interaction graph on the best remaining subgraph (best_subgraph() semantics), then remove the used qubits plus a buffer_hops-neighborhood shell before placing the next copy. Stops at k tiles, when nothing placeable remains, or when the next tile’s score exceeds tile_score_threshold.

Returns tiles ordered best-first; each tile is a positional layout for one circuit copy. For a merged qiskit job, concatenate them tile-major: flat = [q for tile in tiles for q in tile].

n_logical widens each tile beyond the interaction’s edge span: logical qubits without interaction edges (single-qubit-gate-only or idle wires) are assigned best-effort to the best remaining qubits.

With integer k and strict=True, raises ValueError when fewer than k tiles fit; strict=False returns what fits.

Return type:

list[list[int]]

qkan.graph_submodules(model, sample_input, predicate, max_shapes=8, warmup=3)

Wrap every submodule matching predicate with CompiledInference.

Useful for transformer-style models where full-model graph capture fails (e.g. SDPA backends increment an RNG counter even at dropout_p=0). Wrap each MLP block instead — QKAN’s launch-bound cost is concentrated there.

Example:

qkan.graph_submodules(
    gpt2, sample_idx,
    predicate=lambda m: isinstance(m, HQKANMLP),
)

The model is modified in place; replaced submodules wrap the originals so they share parameters (named_parameters is unchanged).

Return type:

Module

qkan.pack_circuit(backend, circuit, k='max', *, tiles=None, profile=None, block_args=(), block_args_batch=None, max_readout_error=None, qubit_error_threshold=None, edge_error_threshold=None, tile_score_threshold=None, buffer_hops=0, strict=True, optimization_level=1)

Pack k copies of a circuit onto disjoint calibrated tiles.

Overloaded over both supported stacks, with the same selection logic (interaction graph -> tile_disjoint() -> copies composed at the tiles’ physical qubits):

• qiskit: pack_circuit(backend, circuit, k) with a QuantumCircuit returns a PackedCircuit — the copies are transpiled pinned to the tiles. profile overrides the calibration snapshot (default DeviceProfile.from_qiskit(backend)) and optimization_level steers the transpilation. Tiles are coupled subgraphs, so the result routes without SWAPs; a routed result (e.g. from a profile whose coupling disagrees with backend) raises RuntimeError.

• CUDA-Q: pack_circuit(profile, kernel, k) with a plain @cudaq.kernel and a DeviceProfile — the kernel’s gates are extracted from the compiled Quake IR (block_args bind runtime arguments) and rebuilt at the tiles’ physical indices with a full-register measurement. A block written over the legacy (q: cudaq.qview, layout: list[int], offset: int) convention is composed by sub-kernel calls exactly as written instead (explicit tiles required); to thread runtime arguments through to sample time, write the whole packed kernel by hand (see the packing guide).

Both stacks require measurement-free inputs: readout belongs to the primitive (qiskit) or is appended automatically (CUDA-Q). Selection thresholds, buffer_hops, k="max", and strict semantics are those of tile_disjoint().

Batched evaluation packs one job per batch: a parameterized qiskit circuit gets per-copy parameters (bind with PackedCircuit.parameter_values()), and a CUDA-Q kernel takes one runtime-argument set per tile via block_args_batch (re-bind the same tiles with PackedCircuit.rebind()). Explicit tiles (hand-picked or from tile_disjoint()) skip the calibration-driven selection on either stack.

Return type:

PackedCircuit

qkan.make_graphed_inference(model, sample_input, warmup=3)

Capture a CUDA graph for model(sample_input) and return a replay fn.

The input to the returned callable must match sample_input in shape, dtype, and device. The output tensor is reused across calls — clone it before issuing the next replay if you need to keep the value around.

Parameters:

• model (Module) – a module in eval() mode. Parameters must not change between capture and replay.

• sample_input (Tensor) – representative input — its shape/dtype/device fixes the capture.

• warmup (int) – warmup forward passes on a side stream before capture. PyTorch recommends >=3 to stabilise cuBLAS/cuDNN selection.

Return type:

Callable[[Tensor], Tensor]

Returns:

A callable fn(x) -> y that replays the captured graph.

qkan.make_graphed_train_step(model, sample_input, sample_target, loss_fn, optimizer, warmup=3)

Capture forward + loss + backward into a CUDA graph for one train step.

The captured region is everything up to (but excluding) optimizer.step: optimizer.zero_grad(set_to_none=False), forward, loss_fn, loss.backward(). The optimizer step is left to the caller because optimisers like Adam contain host-side scalar work that breaks capture.

The returned callable accepts (input, target), copies them into the static buffers, replays the graph, and returns the loss tensor (a view into a static buffer — clone before mutating). After replay, every parameter’s .grad holds the captured-step gradient and optimizer .step() can be called as usual.

Side effects on model and optimizer:

• Every trainable Parameter is re-allocated on the capture stream so the autograd accumulator gets the right stream affinity. References held by the caller (e.g. my_layer.theta) ARE updated; references taken BEFORE calling this function go stale. Re-fetch model.parameters() if you cached anything.

• optimizer.param_groups is rewired to point at the new Parameters; optimizer.state is re-keyed best-effort but state tensors are NOT migrated across streams. Adam users running this on a freshly-built optimizer (no momentum yet) are unaffected.

Constraints:

• sample_input / sample_target fix the shapes and dtypes of subsequent calls; new shapes require a fresh capture.

• All parameters of model must be on CUDA. set_to_none=False inside the graph zeroes grads in-place — never swap .grad for a new tensor outside the graph (e.g. don’t call optimizer.zero_grad(set_to_none=True) between replays).

• loss_fn must produce a scalar tensor.

Parameters:

• model (Module) – a module in train() mode with CUDA parameters.

• sample_input (Tensor) – representative input tensor.

• sample_target (Tensor) – representative target tensor.

• loss_fn (Callable[[Tensor, Tensor], Tensor]) – loss_fn(output, target) -> scalar (e.g. nn.MSELoss()).

• optimizer (Optimizer) – the optimiser whose grads will be zeroed inside the graph. optimizer.step() is NOT captured — call it after each replay.

• warmup (int) – warmup forward/backward passes on a side stream before capture.

Return type:

Callable[[Tensor, Tensor], Tensor]

Returns:

A callable train_step(x, y) -> loss. Loss is a view into a static buffer; clone it if you need to retain past the next replay.

qkan.print0(*s, **kwargs)

qkan.print_banner()

qkan.print_version()

qkan.reshape_optimizer_state(optimizer)

Reshape state tensors to match their parameters’ current shapes.

Walks every per-parameter state entry and, when a tensor’s stored shape doesn’t match its parameter’s current shape but the element count agrees, reshapes it in place. Mismatches with differing numel are left untouched (the optimizer will raise on its next step — that’s the right behaviour for a genuine mismatch).

Returns the number of tensors that were reshaped.

Return type:

int

Solver Module

QKAN solver backends.

Each solver is implemented in its own module: - exact: Pure PyTorch reference implementation - flash: Triton fused kernels - cute: CuTe DSL CUDA fused kernels - cutile: cuTile (NVIDIA Tile Language) fused kernels - cutn: cuQuantum tensor network contraction - qml: PennyLane quantum circuits - qiskit: IBM Quantum backends via Qiskit Runtime - cudaq: NVIDIA CUDA-Q backends (GPU-accelerated simulation or QPU)

class qkan.solver.DeviceProfile(num_qubits, edges=(), readout_error=<factory>, gate_error_1q=<factory>, gate_error_2q=<factory>, t2_us=None, operational=None)

Bases: object

Provider-neutral snapshot of device calibration.

Physical qubit labels run 0..num_qubits-1; devices with non-contiguous or 1-indexed labels simply have holes in the maps, and unreported labels are marked non-operational by the adapters. Coupling edges and the gate_error_2q map are keyed by sorted (low, high) pairs; t2_us is in microseconds. All error values are probabilities in [0, 1]; NaN entries are kept and treated as missing data at scoring time, and missing dict keys mean the device reported nothing for that qubit or edge (common for faulty qubits). Qubits mapped to False in operational are never selected.

num_qubits: int

edges: tuple[tuple[int, int], ...] = ()

readout_error: Mapping[int, float]

gate_error_1q: Mapping[int, float]

gate_error_2q: Mapping[tuple[int, int], float]

t2_us: Mapping[int, float] | None = None

operational: Mapping[int, bool] | None = None

has_calibration()

True when the profile carries any per-qubit calibration data.

Return type:

bool

classmethod from_qiskit(backend, *, refresh=False)

Build a profile from a qiskit backend.

Prefers the BackendV2 Target (structured per-instruction errors); falls back to the legacy properties() API. With refresh=True the backend is asked to re-fetch calibration first — via backend.refresh() on BackendV2 (qiskit-ibm-runtime) or properties(refresh=True) on the legacy path — since both APIs cache their first snapshot.

Return type:

DeviceProfile

classmethod from_braket(device)

Build a profile from an AWS Braket device (experimental).

device is a braket.aws.AwsDevice (or any object exposing the same properties model). Reads the standardized gate-model QPU calibration schema: per-qubit T1/T2, readout fidelity (or IQM’s readout-error entries), and 1Q (simultaneous) randomized-benchmarking fidelity, plus per-pair 2Q gate fidelities. Fidelities are converted to errors as 1 - fidelity; entries whose type name starts with READOUT_ERROR are already errors and are averaged as-is.

Return type:

DeviceProfile

class qkan.solver.PackedCircuit(tiles, circuit=None, isa=None, parameters=(), kernel=None, gates=None, block=None)

Bases: object

A packed job: copies circuit copies pinned to disjoint tiles.

One class serves both stacks — pack_circuit() fills the half that applies, and the placement accessors are shared.

qiskit (estimator-driven): circuit is the merged logical circuit (copy t on logical qubits t*block_qubits ..) and isa its transpilation pinned to layout; map observables per tile with observable() and bind parameter batches with parameter_values() (each copy carries its own renamed parameters, parameters[t] in the source circuit.parameters order) or rebind().

CUDA-Q (counts-driven): kernel applies the block at each tile’s physical indices and (in automatic mode) measures the full register; sample it and read per-tile results with expectation(), X/Y observables through basis_kernel() (hardware-safe) or observable() (simulator observe), and bind argument batches with rebind(). gates holds one extracted gate list per tile and block the source kernel.

tiles: tuple[tuple[int, ...], ...]

circuit: Any = None

isa: Any = None

parameters: tuple[tuple[str, ...], ...] = ()

kernel: Any = None

gates: tuple | None = None

block: Any = None

property layout: tuple[int, ...]

Tile-major flattening of tiles (the pinned initial layout).

property width: int

Register width of the packed CUDA-Q kernel (max index + 1).

property block_qubits: int

Qubits per circuit copy.

property copies: int

Number of packed circuit copies.

physical_qubits(tile)

Physical qubit indices hosting circuit copy tile.

Return type:

list[int]

parameter_values(batch)

Bind one batch of block parameters — copy t gets batch[t].

batch[t][j] binds the j-th parameter (in the source circuit’s circuit.parameters order) of copy t. Returns the name-to-value mapping for an EstimatorV2 PUB alongside isa, so a variational loop packs once and re-binds every step.

Return type:

dict[str, float]

observable(obs, tile=None)

Map a block-level observable onto the packed placement.

Returns the mapped operator for one tile, the list over all copies (None), or their average ("mean") — when every tile runs the same circuit, the mean pools the tiles into one estimate at copies-fold shots.

qiskit packs take a Pauli string (little-endian) or SparsePauliOp over the block’s block_qubits qubits and return the ISA-mapped operator, ready for an EstimatorV2 PUB alongside isa. CUDA-Q packs take a Pauli string in index order (pauli[i] acts on block qubit i) and return a cudaq.spin operator at the tile’s physical indices, for cudaq.observe(packed.observe_kernel(), ...). That is the simulator route — hardware targets that compact idle qubits reject sparse-index observables; use basis_kernel() plus expectation() there.

rebind(batch)

Re-pack the same tiles with a new batch of values.

The per-step call of a variational loop — tile selection is reused. CUDA-Q packs re-extract the block per tile (batch holds one runtime-argument set per tile, as block_args_batch in pack_circuit(); milliseconds per step). qiskit packs return a copy with the parameters bound (batch rows as in parameter_values(); prefer parameter_values in estimator PUBs, which keeps the circuit parametric).

Return type:

PackedCircuit

basis_kernel(bases)

Packed kernel with per-qubit basis rotations before measurement.

bases has one letter per block qubit, position i acting on block qubit i: X and Y append that basis change (H, or RZ(-pi/2) then H) on the corresponding physical qubit of every tile; Z and I leave it untouched. Sample the returned kernel and read Pauli expectations with expectation() — the hardware-safe observable path (sparse-index observe fails on targets that compact idle qubits).

expectation(result, pauli, tile=None)

Pauli expectation values from a sample of the matching basis.

pauli[i] acts on block qubit i; result must come from sampling basis_kernel() for that string (the plain kernel suffices when it only contains Z and I). The value is the Z-parity of the non-identity positions — for one tile, the list over all copies (None), or their average ("mean"): when every tile runs the same circuit, the mean pools the tiles into one estimate at copies-fold shots.

observe_kernel()

Measurement-free packed kernel for cudaq.observe.

observe rejects kernels with measurements, so the sampled kernel (which measures the full register) cannot be used with observable(); this rebuild omits the measurement. Simulator route only — without the full-register measurement, hardware pipelines may compact idle qubits.

class qkan.solver.QKANSolver

Bases: ABC

Common interface for QKAN exact solvers.

All solvers compute postacts of shape (batch, out_dim, in_dim) given:

• x: (batch, in_dim) input tensor

• theta: trainable circuit parameters (shape varies by ansatz)

• preacts_weight, preacts_bias: pre-activation linear params

• reps: number of circuit repetitions

• kwargs: ansatz, group, preacts_trainable, fast_measure, out_dim, dtype

name: str

qkan.solver.best_qubits(backend, n, *, max_readout_error=None, qubit_error_threshold=None, strict=True)

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

\[\mathrm{score}(q) = \mathrm{readout\_error}(q) + \mathrm{sx\_err}(q) + 10^{-4} / \max(T_2(q)\,[\mu s],\, 1)\]

(readout error dominates, \(sx\) error is secondary, short \(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:

Top-n physical qubit indices, sorted by ascending score.

Return type:

list[int]

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,
...     },
... )

qkan.solver.best_subgraph(profile, interaction, *, max_readout_error=None, qubit_error_threshold=None, edge_error_threshold=None, strict=True, call_limit=5000000, max_trials=2500)

Best-calibrated placement of one circuit’s interaction graph.

interaction is either an edge list over logical qubits (2q-gate pairs) or a bare int for an edge-free register (which reduces to rank_qubits()). Returns a positional layout (result[i] hosts logical qubit i), directly usable as a qiskit initial_layout.

Note qiskit’s own transpiler already performs noise-aware VF2 placement at optimization_level 2-3; this function exists for explicit thresholds, custom scoring, level-1 pipelines, and non-qiskit stacks.

Raises ValueError when no placement satisfies the thresholds (strict=True) or returns [] (strict=False). The interaction graph must be connected (tile disconnected pieces via tile_disjoint()) and must embed in the coupling map without routing.

Return type:

list[int]

qkan.solver.cudaq_solver(x, theta, preacts_weight, preacts_bias, reps, **kwargs)

Execute QKAN circuits via NVIDIA CUDA-Q.

Drop-in replacement for torch_exact_solver using CUDA-Q’s GPU-accelerated quantum simulation or QPU backends. Circuits are built as CUDA-Q kernels and expectation values are computed via cudaq.observe().

Supports training via the parameter-shift rule when gradients are needed.

Parameters:

• 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)

• target (str, optional) – CUDA-Q target (e.g., “nvidia”, “nvidia-mqpu”, “qpp-cpu”, “braket”, “iqm”, “quantum_machines”). Set before calling via cudaq.set_target().

• api_key (url / executor /) – Forwarded to cudaq.set_target for REST hardware targets. For “quantum_machines” these select the Qoperator server URL, the executor name, and the X-API-Key credential (the QUANTUM_MACHINES_API_KEY env var also works).

• shots (int, optional) – Number of shots. None for exact statevector expectation. Required for the quantum_machines target.

• expectation_via_sample (bool, optional) – Compute <Z> from sampled bitstring marginals instead of cudaq.observe. Defaults to True on the quantum_machines target, whose REST helper mis-handles multi-term observe.

• initial_layout (optional) – None (default), “auto”, or a list of physical qubit indices for the packed parallel_qubits path on hardware targets. “auto” ranks qubits from calibration data: pass device_profile (a qkan.solver.layout.DeviceProfile) or use the braket target with a machine ARN and amazon-braket-sdk installed. Applied by construction (circuit slot j runs on register index layout[j]); native iqm/oqc/anyon targets preserve these indices, while Braket’s vendor compiler may remap them.

• device_profile (DeviceProfile, optional) – Calibration snapshot used by initial_layout=”auto” and for layout bounds checking (see qkan.solver.layout).

• qubit_error_threshold (max_readout_error /) – Calibration thresholds forwarded to the “auto” qubit ranking.

Return type:

Tensor

Returns:

torch.Tensor

shape: (batch_size, out_dim, in_dim)

qkan.solver.cute_exact_solver(x, theta, preacts_weight, preacts_bias, reps, **kwargs)

CuTe DSL-accelerated exact solver. Drop-in replacement for flash_exact_solver.

Uses fused CuTe CUDA kernels for pz_encoding, rpz_encoding, and real ansatzes. Falls back to torch_exact_solver for unsupported ansatzes.

Key optimizations over Triton/cuTile:

• Shared-memory theta trig caching (eliminates redundant sin/cos across batch)

• __sincosf intrinsics (simultaneous sin+cos per call)

• Warp-shuffle reductions for gradient accumulation

Parameters:

torch_exact_solver. (Same as)

Returns:

(batch_size, out_dim, in_dim)

Return type:

torch.Tensor, shape

qkan.solver.cutile_flash_exact_solver(x, theta, preacts_weight, preacts_bias, reps, **kwargs)

cuTile-accelerated exact solver. Drop-in replacement for flash_exact_solver.

Uses fused cuTile kernels for pz_encoding, rpz_encoding, and real ansatzes. Falls back to torch_exact_solver for unsupported ansatzes.

Parameters:

torch_exact_solver. (Same as)

Returns:

(batch_size, out_dim, in_dim)

Return type:

torch.Tensor, shape

qkan.solver.cutn_solver(x, theta, preacts_weight, preacts_bias, reps, **kwargs)

Tensor network contraction solver using optimal contraction paths.

Expresses the entire quantum circuit as a single tensor network and contracts it using an optimal path from cuQuantum or opt_einsum. The contraction plan is precompiled and cached so repeated forward calls pay no path-finding overhead.

Supports pz_encoding (pz), rpz_encoding (rpz), and real ansatzes. Falls back to torch_exact_solver for unsupported ansatzes or reps > 11.

Parameters:

• x (torch.Tensor) – shape: (batch_size, in_dim)

• theta (torch.Tensor) – shape: (*group, reps+1, n_params_per_gate)

• preacts_weight (torch.Tensor) – shape: (*group, reps)

• preacts_bias (torch.Tensor) – shape: (*group, reps)

• reps (int)

• ansatz (str) – options: “pz_encoding”, “pz”, “rpz_encoding”, “rpz”, “real”

• preacts_trainable (bool)

• fast_measure (bool)

• out_dim (int)

• dtype (torch.dtype)

Return type:

Tensor

Returns:

torch.Tensor

shape: (batch_size, out_dim, in_dim)

qkan.solver.flash_exact_solver(x, theta, preacts_weight, preacts_bias, reps, **kwargs)

Triton-accelerated exact solver. Drop-in replacement for torch_exact_solver.

Uses fused Triton kernels for pz_encoding, rpz_encoding, and real ansatzes. Falls back to torch_exact_solver for unsupported ansatzes.

Parameters:

torch_exact_solver. (Same as)

Returns:

(batch_size, out_dim, in_dim)

Return type:

torch.Tensor, shape

qkan.solver.get_registry()

Return the live registry dict (read-only by convention).

Return type:

dict[str, QKANSolver]

qkan.solver.get_solver(name)

Look up a registered solver by name. Raises KeyError if missing.

Return type:

QKANSolver

qkan.solver.interaction_of(circuit)

Extract a qiskit circuit’s 2-qubit interaction edge list.

Returns one entry per 2-qubit gate (duplicates preserved — they weight the tile scoring). Directives (barriers, measurements) are ignored; gates on three or more qubits raise ValueError.

Return type:

list[tuple[int, int]]

qkan.solver.kernel_interaction_of(block, *args)

Extract a plain CUDA-Q kernel’s 2-qubit interaction edge list.

The CUDA-Q analogue of interaction_of(): one entry per 2-qubit gate, duplicates preserved, measurements ignored, gates on three or more qubits raise ValueError. args bind the kernel’s runtime arguments; closure captures, loops, and list indexing resolve to literal qubit indices (see _quake).

Return type:

list[tuple[int, int]]

qkan.solver.make_base_activation(kind, solver)

Return an nn.Module computing the named activation, matched to the solver.

Return type:

Module

qkan.solver.pack_circuit(backend, circuit, k='max', *, tiles=None, profile=None, block_args=(), block_args_batch=None, max_readout_error=None, qubit_error_threshold=None, edge_error_threshold=None, tile_score_threshold=None, buffer_hops=0, strict=True, optimization_level=1)

Pack k copies of a circuit onto disjoint calibrated tiles.

Overloaded over both supported stacks, with the same selection logic (interaction graph -> tile_disjoint() -> copies composed at the tiles’ physical qubits):

• qiskit: pack_circuit(backend, circuit, k) with a QuantumCircuit returns a PackedCircuit — the copies are transpiled pinned to the tiles. profile overrides the calibration snapshot (default DeviceProfile.from_qiskit(backend)) and optimization_level steers the transpilation. Tiles are coupled subgraphs, so the result routes without SWAPs; a routed result (e.g. from a profile whose coupling disagrees with backend) raises RuntimeError.

• CUDA-Q: pack_circuit(profile, kernel, k) with a plain @cudaq.kernel and a DeviceProfile — the kernel’s gates are extracted from the compiled Quake IR (block_args bind runtime arguments) and rebuilt at the tiles’ physical indices with a full-register measurement. A block written over the legacy (q: cudaq.qview, layout: list[int], offset: int) convention is composed by sub-kernel calls exactly as written instead (explicit tiles required); to thread runtime arguments through to sample time, write the whole packed kernel by hand (see the packing guide).

Both stacks require measurement-free inputs: readout belongs to the primitive (qiskit) or is appended automatically (CUDA-Q). Selection thresholds, buffer_hops, k="max", and strict semantics are those of tile_disjoint().

Batched evaluation packs one job per batch: a parameterized qiskit circuit gets per-copy parameters (bind with PackedCircuit.parameter_values()), and a CUDA-Q kernel takes one runtime-argument set per tile via block_args_batch (re-bind the same tiles with PackedCircuit.rebind()). Explicit tiles (hand-picked or from tile_disjoint()) skip the calibration-driven selection on either stack.

Return type:

PackedCircuit

qkan.solver.qiskit_solver(x, theta, preacts_weight, preacts_bias, reps, **kwargs)

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.

Parameters:

• 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

Return type:

Tensor

Returns:

torch.Tensor

shape: (batch_size, out_dim, in_dim)

qkan.solver.qml_solver(x, theta, reps, **kwargs)

Single-qubit data reuploading circuit using PennyLane.

Parameters:

• x (torch.Tensor) – shape: (batch_size, in_dim)

• theta (torch.Tensor) – shape: (reps, 2)

• reps (int)

• qml_device (str) – default: “default.qubit”

qkan.solver.rank_qubits(profile, n, *, max_readout_error=None, qubit_error_threshold=None, strict=True)

Return the n best-calibrated qubit indices, best first.

Device-independent counterpart of the qiskit solver’s best_qubits: ranks by readout error + 1q gate error (as fidelity-product costs) with a small short-T2 penalty. NaN calibration fails an active threshold and otherwise falls back to pessimistic constants; non-operational qubits are skipped. Returns [] when the profile carries no calibration and no threshold was requested; raises ValueError when a threshold cannot be satisfied (unless strict=False, which returns all usable qubits).

Return type:

list[int]

qkan.solver.register(solver)

Register a solver instance under its name attribute.

Return type:

QKANSolver

qkan.solver.score_layout(profile, interaction, layout, *, gate_counts=None)

Score a candidate layout (lower is better).

layout is positional: layout[i] is the physical qubit hosting logical qubit i. The score is a fidelity-product cost: the sum of count * -ln(1 - error) over every mapped 2q edge plus each logical qubit’s readout/1q-gate cost and a small short-T2 penalty. Logical edges whose physical pair is not coupled cost a large finite penalty (they would require routing). Repeated interaction edges count as their multiplicity. gate_counts optionally overrides weights: int keys are per-logical-qubit 1q+readout multiplicity, sorted-tuple keys are per-logical-edge 2q gate counts.

Return type:

float

qkan.solver.tile_disjoint(profile, interaction, k='max', *, buffer_hops=0, max_readout_error=None, qubit_error_threshold=None, edge_error_threshold=None, tile_score_threshold=None, strict=True, call_limit=50000, max_trials=500, n_logical=None)

Tile the chip with disjoint calibration-aware copies of one circuit.

Greedy peel: repeatedly place the interaction graph on the best remaining subgraph (best_subgraph() semantics), then remove the used qubits plus a buffer_hops-neighborhood shell before placing the next copy. Stops at k tiles, when nothing placeable remains, or when the next tile’s score exceeds tile_score_threshold.

Returns tiles ordered best-first; each tile is a positional layout for one circuit copy. For a merged qiskit job, concatenate them tile-major: flat = [q for tile in tiles for q in tile].

n_logical widens each tile beyond the interaction’s edge span: logical qubits without interaction edges (single-qubit-gate-only or idle wires) are assigned best-effort to the best remaining qubits.

With integer k and strict=True, raises ValueError when fewer than k tiles fit; strict=False returns what fits.

Return type:

list[list[int]]

qkan.solver.torch_exact_solver(x, theta, preacts_weight, preacts_bias, reps, **kwargs)

Single-qubit data reuploading circuit.

Parameters:

• x (torch.Tensor) – shape: (batch_size, in_dim)

• theta (torch.Tensor) – shape: (*group, reps, 2)

• preacts_weight (torch.Tensor) – shape: (*group, reps)

• preacts_bias (torch.Tensor) – shape: (*group, reps)

• reps (int)

• ansatz (str) – options: [“pz_encoding”, “px_encoding”], default: “pz_encoding”

• n_group (int) – number of neurons in a group, default: in_dim of x

Return type:

Tensor

Returns:

torch.Tensor

shape: (batch_size, out_dim, in_dim)

Qiskit Solver

qkan.solver.qiskit_solver(x, theta, preacts_weight, preacts_bias, reps, **kwargs)

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.

Parameters:

• 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

Return type:

Tensor

Returns:

torch.Tensor

shape: (batch_size, out_dim, in_dim)

CUDA-Q Solver

qkan.solver.cudaq_solver(x, theta, preacts_weight, preacts_bias, reps, **kwargs)

Execute QKAN circuits via NVIDIA CUDA-Q.

Drop-in replacement for torch_exact_solver using CUDA-Q’s GPU-accelerated quantum simulation or QPU backends. Circuits are built as CUDA-Q kernels and expectation values are computed via cudaq.observe().

Supports training via the parameter-shift rule when gradients are needed.

Parameters:

• 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)

• target (str, optional) – CUDA-Q target (e.g., “nvidia”, “nvidia-mqpu”, “qpp-cpu”, “braket”, “iqm”, “quantum_machines”). Set before calling via cudaq.set_target().

• api_key (url / executor /) – Forwarded to cudaq.set_target for REST hardware targets. For “quantum_machines” these select the Qoperator server URL, the executor name, and the X-API-Key credential (the QUANTUM_MACHINES_API_KEY env var also works).

• shots (int, optional) – Number of shots. None for exact statevector expectation. Required for the quantum_machines target.

• expectation_via_sample (bool, optional) – Compute <Z> from sampled bitstring marginals instead of cudaq.observe. Defaults to True on the quantum_machines target, whose REST helper mis-handles multi-term observe.

• initial_layout (optional) – None (default), “auto”, or a list of physical qubit indices for the packed parallel_qubits path on hardware targets. “auto” ranks qubits from calibration data: pass device_profile (a qkan.solver.layout.DeviceProfile) or use the braket target with a machine ARN and amazon-braket-sdk installed. Applied by construction (circuit slot j runs on register index layout[j]); native iqm/oqc/anyon targets preserve these indices, while Braket’s vendor compiler may remap them.

• device_profile (DeviceProfile, optional) – Calibration snapshot used by initial_layout=”auto” and for layout bounds checking (see qkan.solver.layout).

• qubit_error_threshold (max_readout_error /) – Calibration thresholds forwarded to the “auto” qubit ranking.

Return type:

Tensor

Returns:

torch.Tensor

shape: (batch_size, out_dim, in_dim)

Error Mitigation

qkan.solver._mitigation._richardson_extrapolate(scale_factors, values)

Lagrange interpolation at x=0 for Zero-Noise Extrapolation.

Given expectation values measured at different noise scale factors, extrapolate to the zero-noise limit.

Parameters:

• scale_factors (list) – noise amplification factors, e.g. [1, 3, 5]

• values (list) – corresponding expectation values at each scale factor

Return type:

float

Returns:

Extrapolated zero-noise expectation value

qkan.solver._mitigation._clip_expvals(expvals)

Clamp expectation values to [-1, 1] (valid range for <Z>).

Return type:

list

qkan.solver._mitigation._apply_mitigation(run_fn, mitigation)

Apply error mitigation to a circuit execution function.

Orchestrates ZNE, multi-shot averaging, and clipping.

Parameters:

• run_fn – callable(scale_factor) -> list[float] of expectation values

• mitigation (dict) – dict with keys “zne”, “n_repeats”, “clip_expvals”

Return type:

list

Returns:

list of mitigated expectation values

Utils Module

Create dataset for regression task.

Adapted from [KindXiaoming/pykan@GitHub 91a2f63](https://github.com/KindXiaoming/pykan/tree/91a2f633be2d435b081ef0ef52a7205c7e7bea9e)

qkan.utils.f_inv(x, y_th)

qkan.utils.f_inv2(x, y_th)

qkan.utils.f_inv3(x, y_th)

qkan.utils.f_inv4(x, y_th)

qkan.utils.f_inv5(x, y_th)

qkan.utils.f_sqrt(x, y_th)

qkan.utils.f_power1d5(x, y_th)

qkan.utils.f_invsqrt(x, y_th)

qkan.utils.f_log(x, y_th)

qkan.utils.f_tan(x, y_th)

qkan.utils.f_arctanh(x, y_th)

qkan.utils.f_arcsin(x, y_th)

qkan.utils.f_arccos(x, y_th)

qkan.utils.f_exp(x, y_th)

qkan.utils.create_dataset(f, n_var=2, f_mode='col', ranges=[-1, 1], train_num=1000, test_num=1000, normalize_input=False, normalize_label=False, device='cpu', seed=0)

Create dataset

Parameters:

• f – function the symbolic formula used to create the synthetic dataset

• ranges – list or np.array; shape (2,) or (n_var, 2) the range of input variables. Default: [-1,1].

• train_num – int the number of training samples. Default: 1000.

• test_num – int the number of test samples. Default: 1000.

• normalize_input – bool If True, apply normalization to inputs. Default: False.

• normalize_label – bool If True, apply normalization to labels. Default: False.

• device – str device. Default: ‘cpu’.

• seed – int random seed. Default: 0.

Returns:

dict

Train/test inputs/labels are dataset[‘train_input’], dataset[‘train_label’], dataset[‘test_input’], dataset[‘test_label’]

Return type:

dataset