Key concepts

This page explains the ideas that the toolkit API is built on. Each section links a rule to the function or error you will meet in practice.

The device runs forward passes only

The NPU computes inference. It doesn’t compute gradients. Every function in the ai module ends in _fprop, short for forward propagation.

To use a model on the NPU, train it on the host first, for example with PyTorch on your CPU or GPU. Then pass the trained weights to the toolkit for inference. The Fourier network example shows this pattern with PyTorch layers that switch backend in eval mode.

Data is bfloat16

Every input array must be a NumPy array of dtype ml_dtypes.bfloat16, and every output comes back as bfloat16. The ml_dtypes package provides the type.

import numpy as np
from ml_dtypes import bfloat16

x = np.random.rand(128, 32).astype(bfloat16)

If you pass float32, the toolkit raises an error before any computation:

TypeError: features must be of dtype <class 'ml_dtypes.bfloat16'> but is float32

bfloat16 keeps the range of float32 but stores only 8 bits of mantissa. That gives about three significant decimal digits. For example, 1.2345678 is stored as 1.234375. Expect small differences when you compare toolkit output with a float32 reference. The introduction example measures them for each operation.

Device lifecycle

The generic module manages devices.

Function	What it does
generic.get_available_npus()	Returns a dictionary that maps each device ID to a serial number.
generic.init_npu(id)	Initialises a device. The docstring says this is optional and not necessary for computation.
generic.release_npu(id)	Releases a device and frees its memory. The docstring says this matters when several users share a device.

Every operation takes an optional device_id argument, which defaults to 0:

import qant_native_computing_toolkit as q

q.generic.init_npu(0)
out = q.ai.relu_fprop(x, device_id=0)
q.generic.release_npu(0)

On the CPU emulator, get_available_npus() returns {0: 'cpu_backend'}, and init_npu, release_npu and device_id have no effect.

Align buffers to page boundaries

The toolkit documentation says the NPU requires all input buffers to be aligned to a 4096-byte page boundary. If a buffer is not aligned, the runtime makes an extra copy, which costs time.

To avoid the copy, align your arrays once with utils.align_ndarray_page_boundary(). It returns the array unchanged if it is already aligned, and an aligned copy otherwise:

features = np.random.rand(128, 32).astype(bfloat16)
weights = np.random.rand(64, 32).astype(bfloat16)

features = q.utils.align_ndarray_page_boundary(features)
weights = q.utils.align_ndarray_page_boundary(weights)

out = q.ai.linear_fprop(features, weights)

Align arrays that you reuse, such as weights, when you load them, not before every call.

Linear weights are transposed

linear_fprop computes features @ weights.T, not features @ weights. The weights have shape (n_out, n_in), the same layout as a PyTorch nn.Linear weight.

Argument	Shape
features	(..., n_in). Extra leading dimensions are batch dimensions.
weights	(n_out, n_in)
Result	(..., n_out)

This means you can pass a PyTorch state dict weight directly after converting it to bfloat16.

Native and built operations

The NPU hardware computes a small set of primitives. In the toolkit’s driver interface, these are:

• element-wise multiplication
• batched matrix-vector multiplication, which covers matrix-vector and matrix-matrix products
• the scaled periodic nonlinearity, a cosine-like function

The native module exposes them as mul_elementwise, linear_fprop and calc_scaled_periodic_nl_fprop.

The ai module builds layers on top of these primitives in the toolkit’s Rust core. In the 2.3.1 source:

Operation	How the Rust core computes it
linear_fprop, conv_fprop, conv_transpose_fprop	Batched matrix-vector multiplication on the device. Convolutions are rearranged into matrix products first.
batchnorm2d_fprop	Normalises on the host, then scales with element-wise multiplication on the device.
calc_kan_layer_fprop	The scaled periodic nonlinearity on the device.
relu_fprop, sigmoid_fprop, softmax_fprop, add_bias_fprop, pooling	Computed in the Rust core on the host, in bfloat16. These don’t call the device primitives.

You call all of them the same way, so this split only matters when you reason about which parts of a model the optical hardware computes.

The cosine-like nonlinearity

calc_scaled_periodic_nl_fprop(u, v) computes tcos(u) * v for each pair of elements. The toolkit describes tcos as a 2π-periodic function with values between -1 and 1, similar to a cosine. On hardware, it is the shape of the optical nonlinearity. On the CPU emulator, it is an exact cosine.

The KAN layer

ai.calc_kan_layer_fprop builds a Kolmogorov-Arnold Network (KAN) layer from tcos. Each output is a sum of cosine-like terms with learnable phases and amplitudes:

y_j = Σ over i and l of  tcos(ks_l · x_i + φ_jil) · ampls_jil

y = q.ai.calc_kan_layer_fprop(features, phis, ampls, ks)

Argument	Shape	Meaning
features	(batch, n_in)	The input, written x in the formula
phis	(n_out, n_in, n_k)	Phase offsets
ampls	(n_out, n_in, n_k)	Amplitudes
ks	(n_k,)	Frequencies, often 1, 2, 3 and so on
Result	(batch, n_out)	The output, written y in the formula

Softmax takes one row at a time

softmax_fprop does not accept batched input. If you pass a 2D array with more than one row, the toolkit raises:

RuntimeError: Rust core error: Toolkit Error: does not support batched input.

Call it once per row, or pass a 1D array. The docstring says that for multidimensional input the first dimension must be 1.

conv_transpose_fprop has the same limit. A 4D input with more than one image in the batch raises the same error on the emulator, so loop over the images instead.

Deprecated multiply functions

The native module still includes mul_i16 and mul_f32 from earlier versions. mul_i16 raises a DeprecationWarning that tells you to use mul_elementwise instead, and a comment in the Rust source says the int16 path will be removed in 2.4. Use mul_elementwise with bfloat16 arrays in new code.

What’s next

• See every operation compared with PyTorch in the introduction example.
• Learn what the emulator does differently from a card in Emulator and hardware.
• Look up error messages in Troubleshooting.