API Reference

This page contains the complete API reference for QKAN.

Core Module

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.

initialize_from_qkan(qkan, x0, sampling=100)

Initialize KAN from a QKAN.

Parameters:

• qkan (QKAN)

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

• sampling (int)

regularization_loss(regularize_activation=1.0, regularize_entropy=1.0)

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', 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, seed=0, **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

Type:

Literal[“qml”, “exact”]

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

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

static clear_ckpts(folder='./model_ckpt')

Clear all checkpoints.

Parameters:

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

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.

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

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

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”

node_attribute()

Get node attribution scores.

Adapted from “pykan”.

property param_size

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_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_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_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

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

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”

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”

to(*args, **kwargs)

Move the model to the specified device.

Parameters:

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

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=False)

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: False

Returns:

dict

Dictionary containing train_loss and test_loss

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, seed=0)

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

Type:

Union[Literal[“qml”, “exact”], 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

_x0

Leave for ResQKANLayer

Type:

Optional[torch.Tensor]

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.

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

property param_size

reset_parameters()

to(*args, **kwargs)

Move the layer to the specified device.

Parameters:

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

property x0

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)

h(is_dagger=False)

Apply Hadamard gate to the state vector.

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

measure_x()

Measure the state vector in the X basis.

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

Return type:

Tensor

measure_y()

Measure the state vector in the Y basis.

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

Return type:

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)

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

s(is_dagger=False)

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

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

state: Tensor

class qkan.TorchGates

Bases: object

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

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 i_gate(shape)

shape: (out_dim, in_dim)

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

Return type:

Tensor

static identity_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 s_gate(shape)

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=torch.complex64)

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)

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

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

Solver Module

QKAN layer simulating solver

This module provides a solver for quantum neural networks using PyTorch or PennyLane

Code author: Jiun-Cheng Jiang (Jim137@GitHub) Contact: [jcjiang@phys.ntu.edu.tw](mailto:jcjiang@phys.ntu.edu.tw)

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

• preacts_weight (torch.Tensor) – shape: (out_dim, in_dim, reps)

• preacts_bias (torch.Tensor) – shape: (out_dim, in_dim, 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)

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