### Execute training loop example

Source: https://github.com/dlyuangod/mlp-kan/blob/main/_autodocs/04-engine.md

A complete example showing how to initialize the model, optimizer, and criterion before running the training loop with train_one_epoch.

```python
from engine import train_one_epoch
from timm.utils import NativeScaler

model = create_model('deit_base_patch16_224_KAN', num_classes=1000)
model.to(device)

optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4)
criterion = DistillationLoss(
    base_criterion=torch.nn.CrossEntropyLoss(),
    teacher_model=teacher,
    distillation_type='soft',
    alpha=0.5,
    tau=4.0
)
loss_scaler = NativeScaler()

train_loader = create_data_loader(...)

for epoch in range(num_epochs):
    metrics = train_one_epoch(
        model=model,
        criterion=criterion,
        data_loader=train_loader,
        optimizer=optimizer,
        device=device,
        epoch=epoch,
        loss_scaler=loss_scaler,
        max_norm=1.0,
        args=args
    )
    print(f"Epoch {epoch} Loss: {metrics['loss']:.4f}")
```

--------------------------------

### kanBlock Usage Example

Source: https://github.com/dlyuangod/mlp-kan/blob/main/_autodocs/01-models-kan.md

Demonstrates initializing the kanBlock and processing a tensor input.

```python
block = kanBlock(dim=768, num_heads=12, hdim_kan=192)
x = torch.randn(2, 196, 768)  # batch=2, patches=196, embed_dim=768
y = block(x)  # shape: [2, 196, 768]
```

--------------------------------

### MoE Usage Example

Source: https://github.com/dlyuangod/mlp-kan/blob/main/_autodocs/09-hi_moe.md

Demonstrates initializing the MoE layer and processing input tensors.

```python
moe = MoE(
    dim=768,
    num_experts=8,
    hidden_dim=3072,
    loss_coef=0.01
)

x = torch.randn(batch_size=32, seq_len=196, dim=768)
output, aux_loss = moe(x)  # output shape: [32, 196, 768]

# Training loss includes auxiliary loss
total_loss = criterion_loss + aux_loss
```

--------------------------------

### Initialize and Use KANLinear

Source: https://github.com/dlyuangod/mlp-kan/blob/main/_autodocs/03-ekan.md

Example demonstrating layer initialization, forward pass, adaptive grid updates, and regularization loss calculation.

```python
layer = KANLinear(
    in_features=64,
    out_features=128,
    grid_size=5,
    spline_order=3,
    scale_spline=1.0
)

# Standard forward pass
x = torch.randn(32, 64)  # batch=32, features=64
y = layer(x)  # shape: [32, 128]

# Adaptive grid update during training
layer.update_grid(x)  # Adjust grid based on input range
y = layer(x)

# Regularization for sparsity
reg_loss = layer.regularization_loss(
    regularize_activation=1.0,
    regularize_entropy=1.0
)
total_loss = criterion_loss + reg_loss
```

--------------------------------

### VisionKAN Model Initialization

Source: https://github.com/dlyuangod/mlp-kan/blob/main/_autodocs/12-configuration.md

Example of initializing a VisionKAN model with specific hyperparameters.

```python
model = create_model(
    'deit_base_patch16_224_KAN',
    num_classes=10,
    drop_rate=0.1,
    drop_path_rate=0.05,
    hdim_kan=192,
    img_size=224
)
```

--------------------------------

### Configure and apply image transformations

Source: https://github.com/dlyuangod/mlp-kan/blob/main/_autodocs/06-datasets.md

Example demonstrating how to configure arguments and apply training and evaluation transforms to an image.

```python
args = argparse.Namespace(
    input_size=224,
    src=False,  # Use RandomResizedCrop
    color_jitter=0.4,
    aa='rand-m9-mstd0.5',
    train_interpolation='bicubic',
    reprob=0.25,
    remode='pixel',
    recount=1,
    eval_crop_ratio=0.875
)

# Training augmentation
train_transform = build_transform(is_train=True, args=args)

# Evaluation preprocessing
eval_transform = build_transform(is_train=False, args=args)

# Usage
from PIL import Image
img = Image.open('image.jpg')
img_train = train_transform(img)  # Augmented
img_eval = eval_transform(img)    # Resized and normalized
```

--------------------------------

### FasterKAN Usage Example

Source: https://github.com/dlyuangod/mlp-kan/blob/main/_autodocs/02-fasterkan.md

Demonstrates initializing a 3-layer network and performing a forward pass with random input data.

```python
# 3-layer network: 64 -> 128 -> 64 -> 32
kan = FasterKAN([64, 128, 64, 32], num_grids=8, denominator=0.33)
x = torch.randn(16, 64)
y = kan(x)  # shape: [16, 32]
```

--------------------------------

### ViT Model Initialization

Source: https://github.com/dlyuangod/mlp-kan/blob/main/_autodocs/12-configuration.md

Example of initializing a ViT model with LayerScale parameters.

```python
model = create_model(
    'deit_base_patch16_LS',
    num_classes=10,
    img_size=224,
    pretrained=True,
    init_scale=1e-4  # LayerScale initialization
)
```

--------------------------------

### FasterKANLayer Usage Example

Source: https://github.com/dlyuangod/mlp-kan/blob/main/_autodocs/02-fasterkan.md

Basic instantiation and forward pass usage of the FasterKANLayer.

```python
layer = FasterKANLayer(input_dim=64, output_dim=128, num_grids=8)
x = torch.randn(32, 64)  # batch=32, features=64
y = layer(x)  # shape: [32, 128]
```

--------------------------------

### Initialize and load datasets

Source: https://github.com/dlyuangod/mlp-kan/blob/main/_autodocs/06-datasets.md

Example usage of build_dataset to load training and validation sets with specific configuration arguments.

```python
from datasets import build_dataset
import argparse

args = argparse.Namespace(
    data_set='CIFAR10',
    data_path='/data',
    input_size=224,
    color_jitter=0.3,
    aa='rand-m9-mstd0.5',
    train_interpolation='bicubic',
    reprob=0.25,
    remode='pixel',
    recount=1,
    eval_crop_ratio=0.875
)

train_dataset, num_classes = build_dataset(is_train=True, args=args)
val_dataset, _ = build_dataset(is_train=False, args=args)

print(f"Training samples: {len(train_dataset)}")
print(f"Classes: {num_classes}")

from torch.utils.data import DataLoader
train_loader = DataLoader(train_dataset, batch_size=128, shuffle=True)
```

--------------------------------

### Initialize and Train KAN Network

Source: https://github.com/dlyuangod/mlp-kan/blob/main/_autodocs/03-ekan.md

Example usage showing network initialization, forward pass, adaptive grid updates, and regularization loss integration.

```python
# 3-layer network: 64 -> 128 -> 64 -> 32
kan = KAN(
    layers_hidden=[64, 128, 64, 32],
    grid_size=5,
    spline_order=3,
    grid_range=[-1, 1]
)

x = torch.randn(16, 64)  # batch=16, features=64

# Standard forward
y = kan(x)  # shape: [16, 32]

# With adaptive grid updates
y = kan(x, update_grid=True)

# Training with regularization
output = kan(x)
sparsity_loss = kan.regularization_loss(
    regularize_activation=1.0,
    regularize_entropy=1.0
)
loss = criterion(output, target) + 0.01 * sparsity_loss
```

--------------------------------

### Common Type Usage Examples

Source: https://github.com/dlyuangod/mlp-kan/blob/main/_autodocs/10-types.md

Summary of common types and their typical usage patterns in the project.

```text
torch.Tensor: torch.randn(batch, channels, height, width)
torch.nn.Module: VisionKAN(...)
Callable[[T], U]: transform(image)
Dict[str, float]: {'loss': 0.5, 'acc1': 90.0}
argparse.Namespace: args = argparse.Namespace(...)
torch.device: torch.device('cuda:0')
torch.optim.Optimizer: torch.optim.AdamW(...)
Iterable[Tuple[T, U]]: DataLoader(...)
```

--------------------------------

### HeirarchicalMoE Usage Example

Source: https://github.com/dlyuangod/mlp-kan/blob/main/_autodocs/09-hi_moe.md

Demonstrates initializing the HeirarchicalMoE module and performing a forward pass with a random input tensor.

```python
hmoe = HeirarchicalMoE(
    dim=768,
    num_experts=(4, 4),  # 4 outer, 4 inner = 16 total
    loss_coef=0.01
)

x = torch.randn(batch_size=32, seq_len=196, dim=768)
output, aux_loss = hmoe(x)  # output shape: [32, 196, 768]

# Total experts: outer * inner = 4 * 4 = 16
# Compute: O(batch * seq * 2 * 2) expert operations
```

--------------------------------

### Initialize and use INatDataset

Source: https://github.com/dlyuangod/mlp-kan/blob/main/_autodocs/06-datasets.md

Example of configuring transformations and initializing the dataset for a specific biological category.

```python
from datasets import INatDataset
import torchvision.transforms as transforms

transform = transforms.Compose([
    transforms.Resize(256),
    transforms.CenterCrop(224),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406],
                        std=[0.229, 0.224, 0.225])
])

dataset = INatDataset(
    root='/path/to/inat',
    train=True,
    year=2018,
    category='family',  # Classify by biological family
    transform=transform
)

# Access dataset
image, label = dataset[0]
print(f"Classes: {dataset.nb_classes}")
```

--------------------------------

### Configure Augmentation Strength

Source: https://github.com/dlyuangod/mlp-kan/blob/main/_autodocs/07-augment.md

Examples of configuring the augmentation pipeline for different intensity levels.

```python
args.src = False           # RandomResizedCrop
args.color_jitter = 0.1    # Minimal color jitter
```

```python
args.src = False           # RandomResizedCrop
args.color_jitter = 0.4    # Strong color shifts
# Plus: random grayscale, solarization, Gaussian blur
```

--------------------------------

### VisionKAN Usage Example

Source: https://github.com/dlyuangod/mlp-kan/blob/main/_autodocs/01-models-kan.md

Initializes a VisionKAN model and performs a forward pass with a random input tensor.

```python
model = VisionKAN(
    patch_size=16, embed_dim=768, depth=12, num_heads=12,
    num_classes=1000, hdim_kan=192
)
x = torch.randn(4, 3, 224, 224)
output = model(x)  # shape: [4, 1000]
```

--------------------------------

### DistilledVisionTransformer Usage Example

Source: https://github.com/dlyuangod/mlp-kan/blob/main/_autodocs/01-models-kan.md

Demonstrates model initialization and conditional output handling based on training or inference mode.

```python
model = DistilledVisionTransformer(
    patch_size=16, embed_dim=768, depth=12, num_heads=12,
    num_classes=1000
)
x = torch.randn(4, 3, 224, 224)
if model.training:
    cls_out, dist_out = model(x)  # Each shape: [4, 1000]
else:
    avg_out = model(x)  # shape: [4, 1000]
```

--------------------------------

### Standard Project Imports and Type Annotations

Source: https://github.com/dlyuangod/mlp-kan/blob/main/_autodocs/10-types.md

Includes essential PyTorch, torchvision, and typing modules. Provides examples of common type annotations for model components and data structures.

```python
from typing import Tuple, Optional, List, Dict, Iterable, Callable

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, Dataset

import torchvision.transforms as transforms
from PIL import Image

import argparse
from argparse import Namespace

# Custom types (implicit)
device: torch.device = torch.device('cuda:0')
model: nn.Module = create_model(...)
optimizer: optim.Optimizer = optim.AdamW(...)
dataloader: Iterable[Tuple[torch.Tensor, torch.Tensor]] = DataLoader(...)
metrics: Dict[str, float] = {...}
```

--------------------------------

### Initialize and use argparse.Namespace

Source: https://github.com/dlyuangod/mlp-kan/blob/main/_autodocs/10-types.md

Demonstrates creating a configuration object and passing it to model and dataset builders.

```python
import argparse

args = argparse.Namespace(
    data_set='CIFAR10',
    model='deit_base_patch16_224_KAN',
    batch_size=128,
    epochs=100,
    learning_rate=5e-4
)

model = create_model(args.model, num_classes=10)
dataset, num_classes = build_dataset(is_train=True, args=args)
```

--------------------------------

### Full Training Pipeline Implementation

Source: https://github.com/dlyuangod/mlp-kan/blob/main/_autodocs/08-kit.md

Demonstrates the end-to-end process of initializing a model, configuring the optimizer, and executing the training and evaluation loops.

```python
from kit import create_model, vit_models
from engine import train_one_epoch, evaluate
import torch
import torch.optim as optim
from torch.utils.data import DataLoader

# Model selection
model = create_model('deit_base_patch16_LS', num_classes=1000)
model.to(device)

# Optimizer and loss
optimizer = optim.AdamW(model.parameters(), lr=5e-4, weight_decay=0.05)
criterion = torch.nn.CrossEntropyLoss()

# Training loop
for epoch in range(num_epochs):
    train_metrics = train_one_epoch(
        model, criterion, train_loader, optimizer,
        device, epoch, loss_scaler, args=args
    )
    
    val_metrics = evaluate(val_loader, model, device)
    print(f"Epoch {epoch}: Train Loss={train_metrics['loss']:.4f}, "
          f"Val Acc@1={val_metrics['acc1']:.2f}%")
```

--------------------------------

### Project Entry Points

Source: https://github.com/dlyuangod/mlp-kan/blob/main/_autodocs/00-index.md

Exposes core model creation, training, and evaluation functions from the package root.

```python
from models_kan import create_model
from engine import train_one_epoch, evaluate

__all__ = ['create_model', 'train_one_epoch', 'evaluate']
```

--------------------------------

### Model Evaluation Example

Source: https://github.com/dlyuangod/mlp-kan/blob/main/_autodocs/04-engine.md

Basic usage pattern for invoking the evaluation function on a validation dataset.

```python
from engine import evaluate

model.eval()
val_loader = create_val_loader(...)

stats = evaluate(
    data_loader=val_loader,
    model=model,
    device=device
)

print(f"Validation Acc@1: {stats['acc1']:.2f}%")
print(f"Validation Acc@5: {stats['acc5']:.2f}%")
print(f"Validation Loss: {stats['loss']:.4f}")
```

--------------------------------

### Build datasets and augmentation pipelines

Source: https://github.com/dlyuangod/mlp-kan/blob/main/_autodocs/00-index.md

Sets up dataset loading and transformation pipelines using argparse configuration.

```python
from datasets import build_dataset, build_transform
import argparse

args = argparse.Namespace(
    data_set='CIFAR10',
    data_path='/path/to/data',
    input_size=224,
    color_jitter=0.4,
    aa='rand-m9-mstd0.5'
)

# Create dataset
dataset, num_classes = build_dataset(is_train=True, args=args)

# Create augmentation pipeline
augmentation = build_transform(is_train=True, args=args)

# DataLoader
from torch.utils.data import DataLoader
loader = DataLoader(dataset, batch_size=128, shuffle=True)
```

--------------------------------

### Integrate MoE into a Transformer block

Source: https://github.com/dlyuangod/mlp-kan/blob/main/_autodocs/09-hi_moe.md

Example implementation of a Transformer block incorporating an MoE layer for the feed-forward network.

```python
from hi_moe import MoE

class TransformerBlockWithMoE(torch.nn.Module):
    def __init__(self, dim, num_heads, num_experts=8):
        super().__init__()
        self.attn = MultiHeadAttention(dim, num_heads)
        self.moe = MoE(dim, num_experts)
        self.norm1 = LayerNorm(dim)
        self.norm2 = LayerNorm(dim)
    
    def forward(self, x):
        # Attention
        x = x + self.attn(self.norm1(x))
        
        # MoE FFN
        moe_out, aux_loss = self.moe(self.norm2(x))
        x = x + moe_out
        
        return x, aux_loss
```

--------------------------------

### Initialize and Run vit_models

Source: https://github.com/dlyuangod/mlp-kan/blob/main/_autodocs/08-kit.md

Demonstrates how to instantiate the model with custom block layers and perform forward passes for classification and feature extraction.

```python
model = vit_models(
    img_size=224,
    patch_size=16,
    embed_dim=768,
    depth=12,
    num_heads=12,
    num_classes=1000,
    block_layers=Layer_scale_init_Block,
    init_scale=1e-4
)

x = torch.randn(4, 3, 224, 224)
output = model(x)  # shape: [4, 1000]

features = model.forward_features(x)  # shape: [4, 768]
```

--------------------------------

### FasterKAN.__init__

Source: https://github.com/dlyuangod/mlp-kan/blob/main/_autodocs/02-fasterkan.md

Initializes a multi-layer FasterKAN network with specified hidden dimensions and spline configuration.

```APIDOC
## FasterKAN.__init__

### Description
Initializes the FasterKAN module, creating a sequence of layers based on the provided hidden dimensions.

### Parameters
- **layers_hidden** (List[int]) - Required - Dimension sequence: [input_dim, hidden1, ..., output_dim]
- **grid_min** (float) - Optional - Basis function grid minimum (default: -2.)
- **grid_max** (float) - Optional - Basis function grid maximum (default: 2.)
- **num_grids** (int) - Optional - Number of spline basis functions (default: 8)
- **denominator** (float) - Optional - Controls smoothness of basis functions (default: 0.33)
- **spline_weight_init_scale** (float) - Optional - Scale factor for spline weight initialization (default: 0.667)
```

--------------------------------

### Use KAN layers directly

Source: https://github.com/dlyuangod/mlp-kan/blob/main/_autodocs/00-index.md

Instantiates a KAN network and demonstrates grid updates and regularization loss calculation.

```python
from ekan import KAN
import torch

# Create KAN network
kan = KAN(
    layers_hidden=[64, 128, 64, 32],
    grid_size=5,
    spline_order=3
)

# Forward pass
x = torch.randn(16, 64)
y = kan(x)  # [16, 32]

# With adaptive grid update
y = kan(x, update_grid=True)

# Regularization loss
reg_loss = kan.regularization_loss(
    regularize_activation=1.0,
    regularize_entropy=1.0
)
```

--------------------------------

### Initialize and use ReflectionalSwitchFunction

Source: https://github.com/dlyuangod/mlp-kan/blob/main/_autodocs/02-fasterkan.md

Shows how to configure the radial basis function module and process input data to generate basis values.

```python
rbf = ReflectionalSwitchFunction(grid_min=-2., grid_max=2., num_grids=8)
x = torch.randn(16, 32)  # batch=16, features=32
bases = rbf(x)  # shape: [16, 32, 8]
```

--------------------------------

### Initialize and use SplineLinear

Source: https://github.com/dlyuangod/mlp-kan/blob/main/_autodocs/02-fasterkan.md

Demonstrates instantiation of the SplineLinear layer and a forward pass with a random input tensor.

```python
layer = SplineLinear(32, 64)
x = torch.randn(16, 32)
y = layer(x)  # shape: [16, 64]
```

--------------------------------

### Initialize MoE and Hierarchical MoE models

Source: https://github.com/dlyuangod/mlp-kan/blob/main/_autodocs/09-hi_moe.md

Configure standard MoE and Hierarchical MoE layers using specific dimensions and expert routing parameters.

```python
# Single-level MoE
moe = MoE(
    dim=768,
    num_experts=8,
    hidden_dim=3072,
    loss_coef=0.01,
    capacity_factor_train=1.25,
    capacity_factor_eval=2.0
)

# Hierarchical MoE
hmoe = HeirarchicalMoE(
    dim=768,
    num_experts=(4, 4),  # 16 total
    loss_coef=0.01
)
```

--------------------------------

### Create and run inference with KAN models

Source: https://github.com/dlyuangod/mlp-kan/blob/main/_autodocs/00-index.md

Initializes a KAN-based model and performs a forward pass with dummy input data.

```python
from models_kan import create_model

# Create KAN model
model = create_model(
    'deit_base_patch16_224_KAN',
    num_classes=1000,
    drop_rate=0.1
)

# Forward pass
import torch
x = torch.randn(4, 3, 224, 224)
logits = model(x)  # [4, 1000]
```

--------------------------------

### Configure Distillation Training

Source: https://github.com/dlyuangod/mlp-kan/blob/main/_autodocs/05-losses.md

Sets up a student-teacher model pair with a distillation loss function and standard optimizer for training.

```python
# Model configuration
model = create_model('deit_tiny_patch16_224_KAN', num_classes=1000)
teacher = create_model('deit_base_patch16_224_ViT', num_classes=1000, pretrained=True)

# Loss configuration
criterion = DistillationLoss(
    base_criterion=torch.nn.CrossEntropyLoss(),
    teacher_model=teacher,
    distillation_type='soft',
    alpha=0.5,
    tau=4.0
)

# Optimizer
optimizer = torch.optim.AdamW(
    model.parameters(),
    lr=5e-4,
    weight_decay=0.05
)

# Training with engine.train_one_epoch
for epoch in range(num_epochs):
    train_metrics = train_one_epoch(
        model=model,
        criterion=criterion,
        data_loader=train_loader,
        optimizer=optimizer,
        device=device,
        epoch=epoch,
        loss_scaler=loss_scaler,
        args=args
    )
```

--------------------------------

### Configure Knowledge Distillation

Source: https://github.com/dlyuangod/mlp-kan/blob/main/_autodocs/12-configuration.md

Sets parameters for teacher-student model distillation, including temperature and weight settings.

```python
args.distillation_type = 'soft'    # 'soft', 'hard', 'none'
args.distillation_alpha = 0.5      # Distillation weight
args.distillation_tau = 4.0        # Temperature for soft targets
args.teacher_model = 'deit_base_patch16_224_ViT'  # Teacher model name
args.teacher_path = None           # Path to teacher checkpoint
```

--------------------------------

### MoE.__init__

Source: https://github.com/dlyuangod/mlp-kan/blob/main/_autodocs/09-hi_moe.md

Initializes the MoE layer with specified dimensions, number of experts, and gating policies.

```APIDOC
## MoE.__init__

### Signature
`MoE(dim: int, num_experts: int = 16, hidden_dim: Optional[int] = None, activation: type = torch.nn.ReLU, second_policy_train: str = 'random', second_policy_eval: str = 'random', second_threshold_train: float = 0.2, second_threshold_eval: float = 0.2, capacity_factor_train: float = 1.25, capacity_factor_eval: float = 2.0, loss_coef: float = 1e-2, experts: Optional[torch.nn.ModuleList] = None)`

### Parameters
- **dim** (int) - Required - Embedding dimension
- **num_experts** (int) - Optional - Number of experts (default: 16)
- **hidden_dim** (int) - Optional - Expert hidden dimension (default: dim * 4)
- **activation** (type) - Optional - Expert activation function (default: nn.ReLU)
- **loss_coef** (float) - Optional - Coefficient for auxiliary loss (default: 0.01)
- **experts** (ModuleList) - Optional - Pre-built expert module; auto-created if None
```

--------------------------------

### Implement Forward Method

Source: https://github.com/dlyuangod/mlp-kan/blob/main/_autodocs/03-ekan.md

Sequential application of KAN layers with optional adaptive grid updates.

```python
def forward(self, x, update_grid=False):
    for layer in self.layers:
        if update_grid:
            layer.update_grid(x)
        x = layer(x)
    return x
```

--------------------------------

### Configure Dataset Settings

Source: https://github.com/dlyuangod/mlp-kan/blob/main/_autodocs/12-configuration.md

Sets the target dataset and its local path. Ensure the dataset name matches supported options like CIFAR10 or IMNET.

```python
args.data_set = 'CIFAR10'        # Dataset: CIFAR10, CIFAR, IMNET, INAT, etc.
args.data_path = '/path/to/data' # Root dataset directory
args.inat_category = 'name'      # iNaturalist target (name, genus, family, etc.)
```

--------------------------------

### Initialize and use FasterKAN

Source: https://github.com/dlyuangod/mlp-kan/blob/main/_autodocs/02-fasterkan.md

Instantiate the FasterKAN layer with specified dimensions and process an input tensor.

```python
# In models_kan.py - MoE_KAN_MLP expert layer
kan = FasterKAN([self.hidden_dim, self.ffn_dim // 2, self.hidden_dim])
output = kan(input_tensor)

# Typical configuration in kanBlock/MoE_KAN_MLP
expert_output = FasterKAN([hidden_dim, ffn_dim // 2, hidden_dim])
```

--------------------------------

### Configure Logging and Checkpointing

Source: https://github.com/dlyuangod/mlp-kan/blob/main/_autodocs/12-configuration.md

Defines output paths and frequency for model training logs and checkpoint saves.

```python
args.output_dir = './output'       # Output directory for checkpoints
args.logging_dir = './logs'        # Logging directory
args.log_freq = 50                 # Logging frequency (batches)
args.save_freq = 1                 # Checkpoint save frequency (epochs)
args.save_model = True             # Save final model
args.resume = None                 # Path to resume checkpoint
```

--------------------------------

### KAN.__init__

Source: https://github.com/dlyuangod/mlp-kan/blob/main/_autodocs/03-ekan.md

Initializes the KAN network with a specified sequence of hidden layers and B-spline hyperparameters.

```APIDOC
## KAN.__init__

### Description
Initializes a multi-layer KAN network by stacking KANLinear layers based on the provided hidden layer dimensions.

### Parameters
- **layers_hidden** (List[int]) - Required - Dimension sequence: [input_dim, hidden1, ..., output_dim]
- **grid_size** (int) - Optional - B-spline grid size per layer (default: 5)
- **spline_order** (int) - Optional - B-spline order (default: 3)
- **scale_noise** (float) - Optional - Noise scale for initialization (default: 0.1)
- **scale_base** (float) - Optional - Base weight initialization scale (default: 1.0)
- **scale_spline** (float) - Optional - Spline weight initialization scale (default: 1.0)
- **base_activation** (type) - Optional - Activation function (default: torch.nn.SiLU)
- **grid_eps** (float) - Optional - Adaptive/uniform grid interpolation (default: 0.02)
- **grid_range** (List[float]) - Optional - Grid range [min, max] (default: [-1, 1])
```

--------------------------------

### Dispatch model creation

Source: https://github.com/dlyuangod/mlp-kan/blob/main/_autodocs/01-models-kan.md

A unified interface that routes to either create_kan or create_ViT based on the model name.

```python
def create_model(model_name: str, **kwargs) -> VisionKAN | VisionTransformer | DistilledVisionTransformer
```

```python
# KAN model
model = create_model('deit_base_patch16_224_KAN', num_classes=1000, drop_rate=0.1)

# Standard ViT
model = create_model('deit_base_patch16_224_ViT', pretrained=True, num_classes=1000)
```

--------------------------------

### Training Script Integration

Source: https://github.com/dlyuangod/mlp-kan/blob/main/_autodocs/06-datasets.md

Standard pattern for building datasets, creating data loaders, and passing them to the training loop.

```python
from datasets import build_dataset, build_transform
import torch.utils.data as data

# Create datasets
train_dataset, num_classes = build_dataset(is_train=True, args=args)
val_dataset, _ = build_dataset(is_train=False, args=args)

# Create data loaders
train_loader = data.DataLoader(
    train_dataset,
    batch_size=args.batch_size,
    shuffle=True,
    num_workers=args.num_workers,
    pin_memory=True
)

val_loader = data.DataLoader(
    val_dataset,
    batch_size=args.batch_size,
    shuffle=False,
    num_workers=args.num_workers,
    pin_memory=True
)

# Pass to training
model = create_model('deit_base_patch16_224_KAN', num_classes=num_classes)

for epoch in range(num_epochs):
    train_one_epoch(model, criterion, train_loader, ...)
    evaluate(val_loader, model, device)
```

--------------------------------

### create_model(model_name, num_classes, drop_rate, pretrained)

Source: https://github.com/dlyuangod/mlp-kan/blob/main/_autodocs/00-index.md

Creates a KAN-based model architecture.

```APIDOC
## create_model

### Description
Instantiates a KAN model by name with specified configuration.

### Parameters
- **model_name** (str) - Required - The name of the model architecture.
- **num_classes** (int) - Optional - Number of output classes.
- **drop_rate** (float) - Optional - Dropout rate for regularization.
- **pretrained** (bool) - Optional - Whether to load pretrained weights.
```

--------------------------------

### create_model(model_name: str, **kwargs)

Source: https://github.com/dlyuangod/mlp-kan/blob/main/_autodocs/12-configuration.md

Initializes a PyTorch model based on the provided model name and configuration arguments.

```APIDOC
## create_model(model_name: str, **kwargs)

### Description
Initializes and returns a torch.nn.Module based on the specified model name and configuration parameters.

### Parameters
- **model_name** (str) - Required - The name of the model architecture to instantiate.
- **num_classes** (int) - Optional - Number of output classes (Default: 1000).
- **pretrained** (bool) - Optional - Whether to load pretrained weights (Default: False).
- **drop_rate** (float) - Optional - Dropout rate on embeddings (Default: 0.0).
- **drop_path_rate** (float) - Optional - Stochastic depth rate (Default: 0.0).
- **img_size** (int) - Optional - Input image size (Default: 224).
- **hdim_kan** (int) - Optional - KAN hidden dimension (Default: 192).
- **batch_size** (int) - Optional - Batch size hint (Default: 16).
- **init_scale** (float) - Optional - LayerScale initialization (ViT-specific).

### Usage Example
```python
model = create_model(
    'deit_base_patch16_224_KAN',
    num_classes=10,
    drop_rate=0.1,
    drop_path_rate=0.05,
    hdim_kan=192,
    img_size=224
)
```
```

--------------------------------

### Initialize and run MoE_KAN_MLP

Source: https://github.com/dlyuangod/mlp-kan/blob/main/_autodocs/01-models-kan.md

Instantiates the MoE_KAN_MLP layer and processes a tensor of hidden states through the gating mechanism.

```python
moe = MoE_KAN_MLP(ffn_dim=768, hidden_dim=384, num_experts=8, top_k=2)
x = torch.randn(16, 12, 384)  # batch_size=16, seq_len=12, hidden_dim=384
output = moe(x)  # shape: [16, 12, 384]
```

--------------------------------

### init_(t)

Source: https://github.com/dlyuangod/mlp-kan/blob/main/_autodocs/09-hi_moe.md

Initializes weight tensors using Xavier uniform initialization.

```APIDOC
## init_(t)

### Description
Applies Xavier uniform initialization to the provided weight tensor.
```

--------------------------------

### Configure PyTorch and CUDA Environment Variables

Source: https://github.com/dlyuangod/mlp-kan/blob/main/_autodocs/12-configuration.md

Set these environment variables in your shell to control GPU visibility, distributed training settings, and CPU threading behavior.

```bash
# CUDA configuration
export CUDA_VISIBLE_DEVICES=0,1,2,3        # GPU selection
export CUDA_LAUNCH_BLOCKING=1               # Synchronous CUDA (for debugging)

# PyTorch distributed
export MASTER_ADDR=localhost                # DDP master address
export MASTER_PORT=29500                    # DDP master port
export RANK=0                               # Process rank
export WORLD_SIZE=1                         # Number of processes

# OpenMP (for CPU threading)
export OMP_NUM_THREADS=8                    # Thread count

# Torch backend selection
export TORCH_HOME=~/.cache/torch            # Cache directory
```

--------------------------------

### Implement Error Prevention Best Practices

Source: https://github.com/dlyuangod/mlp-kan/blob/main/_autodocs/11-errors.md

Use these patterns to validate configurations, verify file paths, check tensor dimensions, monitor training loss, and implement fallback logic for critical operations.

```python
# 1. Validate configuration before training
assert hasattr(args, 'model'), "args.model required"
assert hasattr(args, 'data_set'), "args.data_set required"
assert args.distillation_type in ['none', 'soft', 'hard']

# 2. Check file existence
import os
assert os.path.exists(args.data_path), f"Data path not found: {args.data_path}"

# 3. Validate tensor shapes
assert x.dim() == 2, f"Expected 2D tensor, got {x.dim()}D"
assert x.size(1) == expected_features, f"Feature mismatch"

# 4. Monitor loss during training
if not torch.isfinite(loss):
    raise RuntimeError(f"Loss became {loss}; check data/learning rate")

# 5. Use try-except for critical operations
try:
    dataset, num_classes = build_dataset(is_train=True, args=args)
except FileNotFoundError:
    print("Using fallback dataset configuration")
    args.data_set = 'CIFAR10'
    dataset, num_classes = build_dataset(is_train=True, args=args)
```

--------------------------------

### Dataset Configuration Parameters

Source: https://github.com/dlyuangod/mlp-kan/blob/main/_autodocs/06-datasets.md

Defines the dataset name, root path, and iNaturalist category settings within the args namespace.

```python
args.data_set        # str: Dataset name (CIFAR, CIFAR10, IMNET, etc.)
args.data_path       # str: Root path to dataset
args.inat_category   # str: iNaturalist category (e.g., 'name', 'genus')
```

--------------------------------

### Image Size Configuration

Source: https://github.com/dlyuangod/mlp-kan/blob/main/_autodocs/06-datasets.md

Sets the target image dimensions and the resize ratio used during evaluation.

```python
args.input_size          # int: Target image size (224, 384, etc.)
args.eval_crop_ratio     # float: Resize ratio for eval (0.875 typical)
```

--------------------------------

### Configure Training Loop

Source: https://github.com/dlyuangod/mlp-kan/blob/main/_autodocs/12-configuration.md

Defines core training hyperparameters including batch size, learning rate, and hardware-specific settings.

```python
args.batch_size = 128              # Batch size
args.epochs = 100                  # Number of epochs
args.learning_rate = 5e-4          # Learning rate (AdamW default)
args.weight_decay = 0.05           # Weight decay (L2 regularization)
args.warmup_epochs = 5             # Learning rate warmup
args.min_lr = 1e-5                 # Minimum learning rate
args.num_workers = 4               # DataLoader workers
args.pin_memory = True             # Pin memory for GPU transfer
args.gradient_clip = 1.0           # Max gradient norm for clipping
```

--------------------------------

### MoE Forward Method Implementation

Source: https://github.com/dlyuangod/mlp-kan/blob/main/_autodocs/09-hi_moe.md

Internal logic for routing inputs to experts and combining outputs with auxiliary loss.

```python
def forward(self, inputs):
    # Gate: select top-2 experts
    dispatch, combine, gate_loss = self.gate(inputs)
    
    # Route to experts
    expert_inputs = einsum('bnd,bnec->ebcd', inputs, dispatch)
    
    # Expert computation
    expert_outputs = self.experts(expert_inputs.reshape(e, -1, d))
    
    # Combine outputs
    output = einsum('ebcd,bnec->bnd', expert_outputs, combine)
    
    return output, gate_loss * self.loss_coef
```

--------------------------------

### Project Metadata Configuration

Source: https://github.com/dlyuangod/mlp-kan/blob/main/_autodocs/12-configuration.md

Defines package metadata and dependencies in the pyproject.toml file.

```toml
[tool.poetry]
name = "viskan"
version = "0.1.0"
description = "Viskan"
authors = ["Saaketh Koundinya <three2saki@gmail.com>"]

[tool.poetry.dependencies]
python = "^3.6"
torch = ">=1.13.1"
torchvision = ">=0.8.1"
timm = ">=0.3.2"
```

--------------------------------

### Integrate Augmentation with PyTorch DataLoader

Source: https://github.com/dlyuangod/mlp-kan/blob/main/_autodocs/07-augment.md

Demonstrates how to initialize the augmentation generator and apply it to a custom PyTorch Dataset.

```python
from augment import new_data_aug_generator
from torch.utils.data import DataLoader, Dataset

args = argparse.Namespace(
    input_size=224,
    src=False,
    color_jitter=0.4
)

augmentation = new_data_aug_generator(args)

class ImageDataset(Dataset):
    def __init__(self, image_paths, transform=None):
        self.image_paths = image_paths
        self.transform = transform
    
    def __getitem__(self, idx):
        from PIL import Image
        img = Image.open(self.image_paths[idx])
        if self.transform:
            img = self.transform(img)
        return img

dataset = ImageDataset(
    image_paths=image_list,
    transform=augmentation
)

loader = DataLoader(dataset, batch_size=128, shuffle=True)
for batch in loader:
    # batch shape: [128, 3, 224, 224]
    pass
```

--------------------------------

### Configure iNaturalist dataset parameters

Source: https://github.com/dlyuangod/mlp-kan/blob/main/_autodocs/12-configuration.md

Set path and category-specific parameters for iNaturalist, where the number of classes is determined by the chosen category.

```python
args.data_set = 'INAT'
args.data_path = '/path/to/inat'
args.inat_category = 'name'    # or 'family', 'genus', etc.
args.input_size = 224
args.num_classes = dataset.nb_classes  # Determined by category
```

--------------------------------

### Model Initialization Signature

Source: https://github.com/dlyuangod/mlp-kan/blob/main/_autodocs/12-configuration.md

The function signature for creating a model instance.

```python
create_model(model_name: str, **kwargs) -> torch.nn.Module
```

--------------------------------

### build_dataset(is_train, args)

Source: https://github.com/dlyuangod/mlp-kan/blob/main/_autodocs/06-datasets.md

Creates a vision dataset instance and returns the dataset object along with the number of classes.

```APIDOC
## build_dataset(is_train: bool, args: Namespace)

### Description
Factory function creating vision datasets with standard preprocessing. It returns a tuple containing the PyTorch Dataset instance and the number of classes.

### Parameters
- **is_train** (bool) - Required - Load training split if True, validation if False.
- **args** (Namespace) - Required - Configuration object containing dataset-specific parameters such as 'data_set' and 'data_path'.

### Return Value
- **dataset** (Dataset) - PyTorch Dataset instance ready for DataLoader.
- **num_classes** (int) - Integer number of target classes.

### Example
```python
from datasets import build_dataset
import argparse

args = argparse.Namespace(
    data_set='CIFAR10',
    data_path='/data',
    input_size=224,
    color_jitter=0.3,
    aa='rand-m9-mstd0.5',
    train_interpolation='bicubic',
    reprob=0.25,
    remode='pixel',
    recount=1,
    eval_crop_ratio=0.875
)

train_dataset, num_classes = build_dataset(is_train=True, args=args)
```
```

--------------------------------

### KAN(layers_hidden, grid_size, spline_order)

Source: https://github.com/dlyuangod/mlp-kan/blob/main/_autodocs/00-index.md

Initializes a KAN network layer structure.

```APIDOC
## KAN

### Description
Initializes a Kolmogorov-Arnold Network layer.

### Parameters
- **layers_hidden** (list) - Required - List of hidden layer dimensions.
- **grid_size** (int) - Optional - Size of the spline grid.
- **spline_order** (int) - Optional - Order of the spline functions.
```

--------------------------------

### Apply Solarization to an image

Source: https://github.com/dlyuangod/mlp-kan/blob/main/_autodocs/07-augment.md

Instantiates the Solarization class and applies it to a PIL image.

```python
solar = Solarization(p=0.2)
solarized = solar(pil_image)
```

--------------------------------

### Define argparse.Namespace configuration fields

Source: https://github.com/dlyuangod/mlp-kan/blob/main/_autodocs/10-types.md

Lists the common fields used within the configuration object for training scripts.

```python
# Dataset configuration
data_set: str              # 'CIFAR', 'CIFAR10', 'IMNET', 'INAT', etc.
data_path: str            # Root path to dataset
inat_category: str        # iNaturalist category ('name', 'family', etc.)

# Model configuration
model: str                # Model architecture (e.g., 'deit_base_patch16_224_KAN')
num_classes: int          # Number of classification classes
drop_rate: float          # Dropout rate
drop_path_rate: float     # Stochastic depth drop probability

# Image preprocessing
input_size: int           # Target image size (224, 384)
eval_crop_ratio: float    # Evaluation crop ratio (0.875)
color_jitter: float       # ColorJitter magnitude

# Augmentation
aa: str                   # AutoAugment policy
train_interpolation: str  # Interpolation mode ('bicubic')
reprob: float             # RandAugment probability
remode: str               # RandAugment mode
recount: int              # RandAugment operation count
src: bool                 # Use crop+padding instead of RandomResizedCrop

# Training
batch_size: int           # Batch size
num_workers: int          # DataLoader workers
epochs: int               # Number of training epochs
learning_rate: float      # Learning rate
weight_decay: float       # Weight decay

# Loss and distillation
distillation_type: str    # 'soft', 'hard', 'none'
alpha: float              # Distillation weight
tau: float                # Temperature for soft targets

# Special modes
cosub: bool               # CoSub loss mode
bce_loss: bool            # Use BCE loss with cosub
```

--------------------------------

### Implement Mixture of Experts

Source: https://github.com/dlyuangod/mlp-kan/blob/main/_autodocs/00-index.md

Demonstrates single-level and hierarchical MoE architectures with auxiliary loss integration.

```python
from hi_moe import MoE, HeirarchicalMoE
import torch

# Single-level MoE
moe = MoE(dim=768, num_experts=8, loss_coef=0.01)
x = torch.randn(32, 196, 768)
output, aux_loss = moe(x)

# Hierarchical MoE
hmoe = HeirarchicalMoE(dim=768, num_experts=(4, 4))
output, aux_loss = hmoe(x)

# Include auxiliary loss in training
total_loss = criterion(output, targets) + aux_loss
```

--------------------------------

### Instantiate models with create_model

Source: https://github.com/dlyuangod/mlp-kan/blob/main/_autodocs/08-kit.md

Factory function used to dispatch model creation to registered constructors.

```python
def create_model(model_name: str, **kwargs) -> vit_models
```

```python
model = create_model('deit_tiny_patch16_LS', pretrained=False, num_classes=10)
```

--------------------------------

### Validating Dataset Paths

Source: https://github.com/dlyuangod/mlp-kan/blob/main/_autodocs/11-errors.md

Verify the existence of dataset directories and handle FileNotFoundError exceptions during dataset initialization.

```python
import os
from datasets import build_dataset

# Validate dataset exists
dataset_paths = {
    'CIFAR10': '/afs/crc.nd.edu/user/z/zyuan2/data/cifar10',
    'IMNET': '/scratch365/zyuan2/data/imagenet'
}

for dataset_name, path in dataset_paths.items():
    if not os.path.exists(path):
        print(f"Warning: {dataset_name} not found at {path}")
        # Use alternative dataset or download

# With error handling
try:
    dataset, num_classes = build_dataset(is_train=True, args=args)
except FileNotFoundError as e:
    print(f"Dataset error: {e}")
    # Download or provide correct path
```

--------------------------------

### Handle Grid Update Input Ranges

Source: https://github.com/dlyuangod/mlp-kan/blob/main/_autodocs/11-errors.md

Shows how to avoid degenerate grids by ensuring input data has sufficient variance during update_grid calls.

```python
kan_layer = KANLinear(in_features=64, out_features=128)

# Ensure varied input
x = torch.randn(32, 64)  # Varied distribution
kan_layer.update_grid(x, margin=0.01)

# Avoid constant input
x_const = torch.ones(32, 64)
# Grid update works but becomes single point; not useful
kan_layer.update_grid(x_const, margin=0.01)
```

--------------------------------

### Configure KAN Models

Source: https://github.com/dlyuangod/mlp-kan/blob/main/_autodocs/12-configuration.md

Defines B-spline parameters and initialization settings for KAN model layers.

```python
args.kan_grid_size = 5                      # B-spline grid size
args.kan_spline_order = 3                   # B-spline order (cubic)
args.kan_scale_spline = 1.0                 # Spline weight scale
args.kan_grid_eps = 0.02                    # Adaptive/uniform grid mix
args.kan_scale_noise = 0.1                  # Noise for init
```

--------------------------------

### MoE(dim, num_experts, loss_coef)

Source: https://github.com/dlyuangod/mlp-kan/blob/main/_autodocs/00-index.md

Initializes a Mixture of Experts layer.

```APIDOC
## MoE

### Description
Initializes a single-level Mixture of Experts layer.

### Parameters
- **dim** (int) - Required - Input feature dimension.
- **num_experts** (int) - Required - Number of experts.
- **loss_coef** (float) - Optional - Coefficient for auxiliary loss.
```

--------------------------------

### Configure Image Preprocessing

Source: https://github.com/dlyuangod/mlp-kan/blob/main/_autodocs/12-configuration.md

Defines image dimensions, augmentation policies, and interpolation methods for training inputs.

```python
args.input_size = 224              # Target image size
args.eval_crop_ratio = 0.875       # Eval resize ratio
args.color_jitter = 0.4            # ColorJitter magnitude
args.aa = 'rand-m9-mstd0.5'        # AutoAugment policy
args.train_interpolation = 'bicubic' # Interpolation mode
args.src = False                   # True = crop+pad, False = RandomResizedCrop
```

--------------------------------

### DistillationLoss Forward Method Implementation

Source: https://github.com/dlyuangod/mlp-kan/blob/main/_autodocs/05-losses.md

Internal logic for handling input types, teacher inference, and loss combination.

```python
def forward(self, inputs, outputs, labels):
    # Handle tuple outputs from distilled models
    outputs_kd = None
    if not isinstance(outputs, torch.Tensor):
        outputs, outputs_kd = outputs
    
    # Base loss (student vs labels)
    base_loss = self.base_criterion(outputs, labels)
    
    # No distillation
    if self.distillation_type == 'none':
        return base_loss
    
    # Validate outputs_kd available
    if outputs_kd is None:
        raise ValueError("Distillation enabled but outputs_kd missing")
    
    # Teacher outputs (no gradient)
    with torch.no_grad():
        teacher_outputs = self.teacher_model(inputs)
    
    # Soft/Hard distillation
    if self.distillation_type == 'soft':
        distillation_loss = compute_soft_target_loss(...)
    elif self.distillation_type == 'hard':
        distillation_loss = cross_entropy(outputs_kd, teacher_outputs.argmax(dim=1))
    
    # Combine losses
    loss = base_loss * (1 - self.alpha) + distillation_loss * self.alpha
    return loss
```

--------------------------------

### Create Augmentation Pipeline

Source: https://github.com/dlyuangod/mlp-kan/blob/main/_autodocs/07-augment.md

Constructs a torchvision Compose pipeline using a configuration namespace.

```python
import argparse
from augment import new_data_aug_generator

args = argparse.Namespace(
    input_size=224,
    src=False,
    color_jitter=0.4
)

augmentation = new_data_aug_generator(args)

# Apply to image
from PIL import Image
img = Image.open('image.jpg')
augmented = augmentation(img)  # Returns torch.Tensor [3, 224, 224]
```

--------------------------------

### Create a KAN-based Vision Transformer

Source: https://github.com/dlyuangod/mlp-kan/blob/main/_autodocs/01-models-kan.md

Instantiates a KAN-based model. Distilled variants are currently not supported.

```python
def create_kan(model_name: str, pretrained: bool, **kwargs) -> VisionKAN
```

```python
model = create_kan('deit_base_patch16_224_KAN', pretrained=False, num_classes=10)
```

--------------------------------

### Augmentation Configuration

Source: https://github.com/dlyuangod/mlp-kan/blob/main/_autodocs/06-datasets.md

Configures parameters for color jitter, AutoAugment policies, interpolation, and random erasing settings.

```python
args.color_jitter        # float: ColorJitter magnitude (0.4 typical)
args.aa                  # str: AutoAugment policy ('rand-m9-mstd0.5')
args.train_interpolation # str: Interpolation mode ('bicubic')
args.reprob              # float: RandAugment probability (0.25)
args.remode              # str: RandAugment mode ('pixel', 'const')
args.recount             # int: RandAugment number of operations (1)
```

--------------------------------

### Create a standard Vision Transformer

Source: https://github.com/dlyuangod/mlp-kan/blob/main/_autodocs/01-models-kan.md

Instantiates standard or distilled Vision Transformer models using DeiT weights.

```python
def create_ViT(model_name: str, pretrained: bool, **kwargs) -> VisionTransformer | DistilledVisionTransformer
```

```python
model = create_ViT('deit_base_patch16_224_ViT', pretrained=True, num_classes=10)
```