Mlx-kan

I'm excited to announce the release of MLX-KAN v0.2.5, bringing Kolmogorov-Arnold Networks to the convolutional domain with significant performance improvements!

✨ New Features

🔥 Convolutional KAN Architecture
Introducing a complete convolutional KAN implementation with:

- `KANConv2d`: Core convolutional layer using KAN principles
- `KANConvBlock`: Stackable blocks with optional residual connections
- `KAN_Convolution`: Simplified implementation for backward compatibility
- `KAN_Convolutional_Layer`: Multi-kernel support for legacy code

🧩 Flexible Input Handling
- Support for various input shapes (2D, 3D, 4D tensors)
- Automatic reshaping for different data formats
- Proper padding and patch extraction

🚀 Performance Improvements

- **8x Faster Processing**: Significant optimizations for both inference and training
- Streamlined patch extraction and processing
- Efficient memory usage with improved tensor operations

🛠️ Technical Enhancements

- Proper handling of kernel sizes, strides, padding and dilation
- Support for grouped convolutions
- Regularization loss propagation through convolutional layers
- Helper functions for backward compatibility

🔮 What's Next?

This release lays the groundwork for building more complex KAN-based vision models. Try building a KAN-based CNN and let me know your results!

---

Get started with mlx-kan v0.2.5 today:
bash
pip install kan-mlx==0.2.5

This Release v0.2.0 covers big changes!!!

You can now easily train, save and load your KAN model easily with you r own data. This is currently only supported with the basic KAN model, more support for different model and architectures will come later, also more training algorythms, currently only the SimpleTrainer.

Here is an example using the mnist:

python
Import necessary libraries
import mlx.core as mx Import mlx core module for array manipulation

Import KAN model and arguments
from mlx_kan.kan import KAN
from mlx_kan.kan.args import ModelArgs

Import SimpleTrainer and training arguments
from mlx_kan.trainer.simpletrainer import SimpleTrainer
from mlx_kan.trainer.trainer_args import TrainArgs

Import MNIST dataset loader
import mlx_kan.quick_scripts.mnist as mnist

Define the model parameters
num_layers = 2 Number of layers in the model
in_features = 28 Input feature dimension (e.g., image width for MNIST)
out_features = 28 Output feature dimension (e.g., image height for MNIST)
hidden_dim = 64 Dimension of hidden layers
num_classes = 10 Number of output classes for classification (e.g., digits 0-9)

Initialize the KAN model with specified architecture
kan_model = KAN(
layers_hidden=[in_features * out_features] + [hidden_dim] * (num_layers - 1) + [num_classes], Define layers: input layer, hidden layers, output layer
args=ModelArgs Pass model arguments
)

Load the MNIST dataset
train_images, train_labels, test_images, test_labels = map(mx.array, getattr(mnist, "mnist")()) Convert dataset to mlx arrays

Set training arguments
TrainArgs.max_steps = 1000 Maximum number of training steps

Initialize and run the SimpleTrainer
SimpleTrainer(
model=kan_model, Model to be trained
args=TrainArgs, Training arguments
train_set=(train_images, train_labels), Training dataset
validation_set=(test_images, test_labels), Validation dataset (using test set for validation)
test_set=(test_images, test_labels), Testing dataset
validation_interval=1000, Interval for validation
logging_interval=10 Interval for logging
)


More examples are in the folder `examples`

Adding architectures into the mix.

Added architectures

`SmallKANMLP` Class

The `SmallKANMLP` class consists of two `KANLinear` layers. It is designed for small-scale models.

`MiddleKANMLP` Class

The `MiddleKANMLP` class consists of three `KANLinear` layers. It is designed for medium-scale models.

`BigKANMLP` Class

The `BigKANMLP` class consists of four `KANLinear` layers. It is designed for large-scale models.

`LlamaKANMLP` Class

The `LlamaKANMLP` class consists of three `KANLinear` layers configured in a the same manner Llama's MLP layer is configured. It is designed for models requiring a unique layer arrangement.


How to access?

python
from mlx_kan.kan.architectures.KANMLP import LlamaKANMLP, SmallKANMLP, MiddleKANMLP, BigKANMLP

This release optimizes the model and training code.

Added:
- progress bar for the steps in every epoch
- gradient clipping in the quick training code
- more clean up's

This is the first Package release:

install the Package:

sh
pip install mlx-kan


Example usage in Python:

python
from kan_mlx.kan import KAN

Initialize and use KAN
kan_model = KAN()


If you want some quck playing around:

sh
python -m mlx-kan.quick_scripts.quick_train --help