Train a Fourier network and run it on the NPU

This example compares a conventional multilayer perceptron (MLP) with a Fourier network on the same task: learning a one-dimensional function. The Fourier network uses layers whose building block is a cosine, which the NPU evaluates natively.

Objectives

• Build PyTorch layers that train on the CPU and run inference through the toolkit.
• Train an MLP and a Fourier network with the same budget of steps.
• Run both models through the toolkit and compare their accuracy.

Before you begin

1. Install the toolkit. The CPU emulator works for this example, so you don’t need hardware.

2. Check that Python can import the toolkit:

   python -c "import qant_native_computing_toolkit as q; print(q.generic.get_available_npus())"

   On the emulator this prints {0: 'cpu_backend'}. On a server with cards it lists each NPU.

3. Get the examples repository:

   git clone --recurse-submodules https://github.com/Q-ANT-GmbH/qant_native_computing_toolkit_examples.git
   cd qant_native_computing_toolkit_examples/python/fourier_network

4. Install the example’s dependencies:

   pip install -r requirements.txt

   The file pins torch==2.11.0. A newer PyTorch also works. We ran it with 2.14.1.

Run the example

1. Start Jupyter and open function_learning_on_npu.ipynb:

   jupyter notebook function_learning_on_npu.ipynb

2. Run all cells. The notebook trains both models for 6,000 steps each, then evaluates them through the toolkit.

To run it without opening a browser, execute the notebook from the command line:

jupyter nbconvert --to notebook --execute function_learning_on_npu.ipynb --output executed.ipynb

How it works

The task

Both models learn the same target on the interval [-2, 2]. It mixes a slow cosine with a faster sine:

def func_1d(x: np.ndarray) -> np.ndarray:
    return 0.5 * np.cos(2 * np.pi * x) + 0.25 * np.sin(2 * np.pi * 4 * x)

Layers that switch backends

The file qant_layers.py defines two PyTorch layers, QLinear and QFourier. Each layer has two paths:

• In training mode, it uses ordinary PyTorch maths, so autograd computes gradients on your CPU.
• In eval mode, it converts inputs and weights to bfloat16 and calls the toolkit.

This is the eval path of QFourier. One call computes the whole layer:

input_np = input.detach().numpy().astype(bfloat16)
amplitude_np = self.amplitude.data.detach().numpy().astype(bfloat16)
phase_np = self.phase.data.detach().numpy().astype(bfloat16)
k_np = self.get_buffer("k").numpy().astype(bfloat16)

output_np = qant.ai.calc_kan_layer_fprop(input_np, phase_np, amplitude_np, k_np)

The bias is added on the host after the call returns.

The two models

Model	Layer widths	Grid	Parameters
MLP	[1, 32, 32, 1] with ReLU	none	1,153
Fourier network	[1, 4, 1]	[4, 4]	69

A Fourier layer has in × out × grid × 2 + out parameters: one amplitude and one phase for each frequency on each connection, plus a bias per output.

QFourier also adds Gaussian noise with a standard deviation of 0.04 to its output during training. This makes the learned weights more tolerant of the analog noise a physical device adds.

Inference through the toolkit

Switch both models to eval mode, open the device, run, and release it:

mlp_model.eval()
fourier_model.eval()

qant.generic.init_npu(0)
with torch.no_grad():
    y_pred_mlp = mlp_model(x_tensor).numpy().squeeze()
    y_pred_fourier = fourier_model(x_tensor).numpy().squeeze()
qant.generic.release_npu(0)

Expected output

These are the values printed when we ran the notebook on the CPU emulator. Your numbers can differ slightly between runs and machines.

Measure	MLP	Fourier network
Parameters	1,153	69
Final training loss	0.0244	0.0036
Test MSE through the toolkit	0.0237	0.0009

The notebook ends with a challenge section that fits a harder function. On the emulator it printed a final training loss of 0.04280 and a test MSE of 0.04065.

Run the tests

The example includes a test that executes the notebook and checks the result:

python -m pytest -v test_function_learning.py

On the emulator the test passed in 25.1 seconds.

What’s next

• See the same layer built from scratch in the live cosine example, with Python that runs in your browser.
• Design a larger Fourier network in the playground.
• Read why the bfloat16 format sets a floor on accuracy in the Lab.