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)[source]

Bases: Module

forward(x)[source]

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.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)[source]

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)[source]

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)[source]

initialize_from_qkan(qkan, x0, sampling=100)[source]

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, **kwargs)[source]

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)[source]

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)[source]

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)[source]

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

Parameters:: another_model (QKAN) – Another model to initialize from

initialize_parameters()[source]: Reinitialize parameters of all QKANLayer layers in-place.

xavier_init()[source]: Apply Xavier normal initialization to all QKANLayer layers.

refine(new_reps)[source]

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)[source]

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)[source]

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)[source]

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()[source]

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)[source]

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

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)[source]

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)[source]

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)[source]

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)[source]

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)[source]

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)[source]

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')[source]

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')[source]

Clear all checkpoints.

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

save_ckpt(name, folder='./model_ckpt')[source]

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')[source]

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.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)[source]

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()[source]

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()[source]

Apply Xavier normal initialization to theta and preacts.

Applies nn.init.xavier_normal_ in-place to self.theta. When self.preact_init is set, also applies it to self.preacts_weight and self.preacts_bias.

to(*args, **kwargs)[source]

Move the layer to the specified device.

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

property param_size

property x0

forward(x)[source]

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.

reset_parameters()[source]

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)[source]

get_subset(in_id, out_id)[source]

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.StateVector(batch_size, out_dim, in_dim, device='cpu', dtype=torch.complex64)[source]

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)[source]

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)[source]

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)[source]

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)[source]

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

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

h(is_dagger=False)[source]

Apply Hadamard gate to the state vector.

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

x()[source]: Apply Pauli-X gate to the state vector.

z()[source]: Apply Pauli-Z gate to the state vector.

rx(theta, is_dagger=False)[source]

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)[source]

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)[source]

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[source]

Bases: object

static identity_gate(shape)[source]

shape: (out_dim, in_dim)