Skip to main content

Photonic KAN: Vendor-agnostic PyTorch + KAN framework for photonic AI hardware

Project description

PhotoKAN

Photonic Kolmogorov-Arnold Networks — vendor-agnostic PyTorch framework for photonic AI hardware.

Python PyTorch PyPI


Why PhotoKAN?

Standard MLPs use fixed activations on nodes and linear weights on edges. KANs invert this: learnable nonlinear functions sit on the edges, summed at nodes.

Photonic hardware is physically structured around edge-level nonlinear transforms — light through a waveguide produces exactly the kind of parametric nonlinear function a KAN edge needs. This is not an analogy; it is a direct structural match.

Published benchmarks show:

  • 43% fewer parameters vs equivalent MLPs
  • 46% fewer operations vs equivalent MLPs
  • 30x energy efficiency gain on photonic NPU vs CMOS

Supported Hardware

Vendor Technology SDK Install
Q.ANT Thin-Film Lithium Niobate (TFLN) Q.PAL pip install photokan[qant]
Lightmatter Silicon Photonics Lightmatter SDK pip install photokan[lightmatter]
Salience Labs III-V Photonics (InP) Salience SDK pip install photokan[salience]
CPU Sim Physics-accurate noise model Built-in pip install photokan

Installation

pip install photokan                        # CPU simulation (no hardware required)
pip install photokan[qant]                  # + Q.ANT NPU support
pip install photokan[lightmatter]            # + Lightmatter support
pip install photokan[salience]              # + Salience Labs support
pip install photokan[llm]                   # + HuggingFace / PEFT integration
pip install photokan[onnx]                  # + ONNX export
pip install photokan[dev]                   # + development tools

Quick Start

import torch
import photokan as pk

# Works on CPU sim if no hardware present — no code changes needed
model = pk.PhotoKAN(
    layer_sizes=[4, 16, 16, 1],
    activation='sine',   # 'sine' | 'fourier' | 'spline' | 'relu'
    backend='auto',      # auto-detects photonic hardware, falls back to CPU
)

x = torch.randn(32, 4)
y = model(x)

# Standard PyTorch training
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
loss = torch.nn.MSELoss()(y, torch.randn(32, 1))
loss.backward()
optimizer.step()

# Check what hardware is available
print(pk.available_backends())
# -> {'cpu': True, 'cuda': False, 'qant': False, 'lightmatter': False, 'salience': False}

# List registered vendors
print(pk.all_vendor_names())
# -> ['qant', 'lightmatter', 'salience']

# Get vendor-specific noise profiles
print(pk.get_noise_config('qant', 'npu1'))
# -> {'snr_db': 14.0, 'bit_depth': 6, 'phase_noise_rad': 0.02, 'technology': 'tfln', 'enabled': True}

Vendor-Specific Dispatch

Target a specific vendor or let the framework auto-detect:

model = pk.PhotoKAN(
    layer_sizes=[4, 16, 1],
    activation='sine',
    backend='qant',        # explicitly use Q.ANT hardware
)

# Or simulate a specific vendor's noise characteristics on CPU
model = pk.PhotoKAN(
    layer_sizes=[4, 16, 1],
    activation='sine',
    backend='cpu',
    noise_config=pk.get_noise_config('lightmatter', 'envise1'),
)

Activation Variants

Name Formula Best for Photonic native
sine sum w*sin(f*x + p) Periodic targets, photonic deployment Yes
fourier a0 + sum [a*cos + b*sin] Multi-frequency signals Yes
spline B-spline basis Non-periodic, high precision Via LUT
relu sum w*ReLU(a*x + b) Edge inference, speed Yes

Photonic Noise Simulation

Test accuracy against realistic hardware impairments before deploying:

sim = pk.PhotonicSimulator()

# Use vendor-specific noise profile
sim.set_hardware_profile('npu2')   # Q.ANT NPU2: SNR=16dB, 8-bit

results = sim.sweep_snr(model, x_test, y_test,
                         snr_range=[8, 10, 12, 14, 16, 20])
sim.plot_snr_accuracy(results)

Noise Profiles by Vendor

Q.ANT (TFLN)

Profile SNR Bit Depth Phase Noise
npu1 14 dB 6-bit 0.02 rad
npu2 16 dB 8-bit 0.01 rad
ideal 60 dB 16-bit 0.0 rad

Lightmatter (Silicon Photonics)

Profile SNR Bit Depth Thermal Drift
envise1 12 dB 5-bit 0.005
mars1 15 dB 7-bit 0.002
ideal 60 dB 16-bit 0.0

Salience Labs (III-V/InP)

Profile SNR Bit Depth Ring Crosstalk
mr100 18 dB 8-bit 0.003
mr200 22 dB 10-bit 0.001
ideal 60 dB 16-bit 0.0

Energy Estimation

Estimate energy consumption with published photonic efficiency numbers:

from photokan.utils import estimate_model_energy

reports = estimate_model_energy(model, batch_size=64)
# Layer 0: 8.214 uJ (CMOS) -> 0.267 uJ (Photonic), 30.8x better
# Layer 1: ...

Convolutional KAN

Use PhotoConvKAN for image and spatial workloads:

model = pk.PhotoConvKAN(
    in_channels=1, out_channels=16,
    kernel_size=3, activation='sine',
)
x = torch.randn(8, 1, 28, 28)
y = model(x)  # [8, 16, 28, 28]

LLM Integration

Replace MLP layers in HuggingFace transformers with PhotoKAN using LoRA-style adapters:

from photokan.llm import add_photo_lora, PhotoKANAttention

# Wrap a transformer's MLP layers
model = add_photo_lora(model, target_modules=["mlp"], n_basis=4)

AOT Compilation

Compile models to photonic deployment bundles (.npu):

compiler = pk.PhotonicCompiler()
program = compiler.compile(model, './my_model.npu')

# CPU validation (LUT interpreter)
y = program.run(x, backend='cpu')

# Hardware inference
y = program.run(x, backend='qant')      # Q.ANT
y = program.run(x, backend='lightmatter')  # Lightmatter

# Benchmark latency
stats = program.benchmark(x)
# -> {'mean_ms': 0.42, 'throughput_samples_per_sec': 76190}

Architecture

User PyTorch model
  |- PhotoKAN / PhotoKANLayer / PhotoConvKAN  (nn.Module, drop-in)
       |- EdgeActivations (sine / fourier / spline / relu)
            |- SimBackend -> CPU physics simulation (vendor noise profiles)
            |- Vendor dispatch -> Q.ANT / Lightmatter / Salience via pluggable backends
  |- PhotonicCompiler
       |- LUTCompiler -> int8 quantised lookup tables
       |- GraphBuilder -> photonic execution graph
       |- PhotonicProgram -> run on hardware or CPU LUT interpreter
  |- PhotonicSimulator -> SNR sweeps, transfer functions
  |- LLM integration (LoRA adapters, attention)

Adding a New Vendor

PhotoKAN uses a pluggable backend architecture. To add support for a new photonic vendor:

  1. Create photokan/backends/<vendor>/__init__.py and backend.py
  2. Implement the PhotonicBackend ABC:
from photokan.backends.base import PhotonicBackend
from photokan.backends.registry import register_vendor

class MyBackend(PhotonicBackend):
    @staticmethod
    def name() -> str: ...
    @staticmethod
    def is_available() -> bool: ...
    @staticmethod
    def execute(x, activation, op_type) -> torch.Tensor: ...
    @staticmethod
    def compute_gradient(grad_output, x, activation, op_type) -> tuple: ...
    @staticmethod
    def noise_profiles() -> dict[str, dict]: ...

register_vendor(MyBackend)
  1. Add the vendor to _discover_vendors() in registry.py

Project Status

Phase Features Status
Phase 1 Activations, SimBackend, Layers, Energy, Profiler Done
Phase 2 LUT compiler, execution graph, ONNX export, ConvKAN Done
Phase 3 Vendor-agnostic backends (Q.ANT, Lightmatter, Salience) Done
Phase 4 Hardware dispatch, LLM fine-tuning, arXiv paper Upcoming

Running Tests

pip install -e ".[dev]"
pytest                          # 136+ tests
pytest -k "not slow"            # skip slow tests
pytest --cov=photokan           # with coverage

References


PhotoKAN v0.4.0 -- Build the bridge. Light does the math.

Project details


Download files

Download the file for your platform. If you're not sure which to choose, learn more about installing packages.

Source Distribution

photokan-0.4.1.tar.gz (48.0 kB view details)

Uploaded Source

Built Distribution

If you're not sure about the file name format, learn more about wheel file names.

photokan-0.4.1-py3-none-any.whl (62.2 kB view details)

Uploaded Python 3

File details

Details for the file photokan-0.4.1.tar.gz.

File metadata

  • Download URL: photokan-0.4.1.tar.gz
  • Upload date:
  • Size: 48.0 kB
  • Tags: Source
  • Uploaded using Trusted Publishing? No
  • Uploaded via: twine/6.2.0 CPython/3.12.11

File hashes

Hashes for photokan-0.4.1.tar.gz
Algorithm Hash digest
SHA256 6de7283334cd5613765bbbf3859721a2a5e314c2984052358f67d7d0afafba79
MD5 bab5f43dcb55b15cb157b4ac12b0c76e
BLAKE2b-256 a5630f7313baaf88c25330623426186bf13fa55139406ead5fef3db414deb5a2

See more details on using hashes here.

File details

Details for the file photokan-0.4.1-py3-none-any.whl.

File metadata

  • Download URL: photokan-0.4.1-py3-none-any.whl
  • Upload date:
  • Size: 62.2 kB
  • Tags: Python 3
  • Uploaded using Trusted Publishing? No
  • Uploaded via: twine/6.2.0 CPython/3.12.11

File hashes

Hashes for photokan-0.4.1-py3-none-any.whl
Algorithm Hash digest
SHA256 c7e48e2f13697ef7bd9fb14b3e24bdc0e10cde7f24fd0e08e31ed7b279775897
MD5 8a7765a50809f611ba47ce5988e8dae9
BLAKE2b-256 8f7f202de18bd8b0624f2333a6bcfe1da20160b27bb5424a273fc1949f2bbb54

See more details on using hashes here.

Supported by

AWS Cloud computing and Security Sponsor Datadog Monitoring Depot Continuous Integration Fastly CDN Google Download Analytics Pingdom Monitoring Sentry Error logging StatusPage Status page