chore: refactor PLIF and LIF modules for improved readability and performance

- Enhanced the PLIF and LIF classes by restructuring the initialization parameters for better clarity.
- Updated the forward methods in both classes to improve code readability and maintainability.
- Added support for vectorized decay and threshold parameters in PLIF.
- Improved the handling of surrogate gradients in both LIF and PLIF.
- Refactored the attention mechanism in SpikingSelfAttention to streamline operations.
- Updated the Monitor utility to enhance voltage trace monitoring capabilities.
- Added comprehensive tests for new features and ensured backward compatibility.
- Cleaned up code formatting across multiple files for consistency.

Author: Khushiyant

.github/workflows/docs.yml

@@ -38,6 +38,10 @@ jobs:
         run: |
           uv sync --extra docs
 
+      - name: Ruff Check
+        run: |
+          uv run ruff check .
+
       - name: Build Sphinx Documentation
         run: |
           # Generate RST files if they are not already committed

benchmarks/benchmark_lif.py

@@ -1,14 +1,14 @@
 import torch
-import triton
 import time
 import sys
 import os
 
 # Add src to path so we can import tether
-sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '../src')))
+sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "../src")))
 
 from tether.functional.lif import LIFSubFunction
 
+
 def lif_pytorch(x_seq, v_init, decay, threshold):
     """
     PyTorch reference implementation of LIF.
@@ -31,25 +31,26 @@ def lif_pytorch(x_seq, v_init, decay, threshold):
     """
     # x_seq: (Time, Neurons)
     # v_init: (Neurons)
-    
+
     n_steps, n_neurons = x_seq.shape
     v = v_init.clone()
     spikes_list = []
-    
+
     # We use a simple loop as LIF is recurrent
     # This simulates "Vanilla PyTorch" without custom CUDA kernels or JIT
     for t in range(n_steps):
         x = x_seq[t]
         v = v * decay + x
-        
+
         spike = (v >= threshold).float()
         spikes_list.append(spike)
-        
+
         # Hard reset
         v = v * (1.0 - spike)
-        
+
     return torch.stack(spikes_list)
 
+
 def benchmark():
     """
     Benchmark PyTorch vs Triton implementations of LIF.
@@ -64,10 +65,10 @@ def benchmark():
     # Dimensions
     # Simulating a reasonable layer size for a Transformer
     batch_size = 32
-    seq_len = 2048 # Longer context to emphasize the loop overhead vs kernel
-    dim = 768      # Standard BERT-base dimension
+    seq_len = 2048  # Longer context to emphasize the loop overhead vs kernel
+    dim = 768  # Standard BERT-base dimension
     n_neurons = batch_size * dim
-    
+
     # Inputs
     x_seq = torch.randn(seq_len, n_neurons, device=device)
     v_init = torch.zeros(n_neurons, device=device)
@@ -81,7 +82,7 @@ def benchmark():
         with torch.no_grad():
             _ = lif_pytorch(x_seq, v_init, decay, threshold)
             LIFSubFunction.apply(x_seq, v_init, decay, threshold, alpha, 0)
-    
+
     # Benchmark PyTorch
     torch.cuda.synchronize()
     start_time = time.time()
@@ -103,8 +104,9 @@ def benchmark():
     torch.cuda.synchronize()
     triton_time = (time.time() - start_time) / iterations
     print(f"Triton Time:  {triton_time * 1000:.3f} ms")
-    
+
     print(f"Speedup: {pytorch_time / triton_time:.2f}x")
 
+
 if __name__ == "__main__":
     benchmark()

benchmarks/benchmark_plif.py

@@ -1,15 +1,15 @@
 import torch
-import triton
 import time
 import sys
 import os
 
 # Add src to path
-sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '../src')))
+sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "../src")))
 
 from tether.functional.lif import LIFSubFunction
 from tether.functional.plif import PLIFSubFunction
 
+
 def benchmark_plif():
     """
     Benchmark LIF vs PLIF (Triton vs Triton).
@@ -25,15 +25,15 @@ def benchmark_plif():
     seq_len = 2048
     dim = 768
     n_neurons = batch_size * dim
-    
+
     x_seq = torch.randn(seq_len, n_neurons, device=device)
     v_init = torch.zeros(n_neurons, device=device)
-    
+
     # LIF Params (Scalar)
     decay_scalar = torch.tensor(0.9, device=device)
     threshold_scalar = torch.tensor(1.0, device=device)
     alpha = torch.tensor(2.0, device=device)
-    
+
     # PLIF Params (Vector)
     decay_vector = torch.full((n_neurons,), 0.9, device=device)
     threshold_vector = torch.full((n_neurons,), 1.0, device=device)
@@ -42,16 +42,22 @@ def benchmark_plif():
     print("Warming up...")
     for _ in range(10):
         with torch.no_grad():
-            LIFSubFunction.apply(x_seq, v_init, decay_scalar, threshold_scalar, alpha, 0)
-            PLIFSubFunction.apply(x_seq, v_init, decay_vector, threshold_vector, alpha, 0)
-    
+            LIFSubFunction.apply(
+                x_seq, v_init, decay_scalar, threshold_scalar, alpha, 0
+            )
+            PLIFSubFunction.apply(
+                x_seq, v_init, decay_vector, threshold_vector, alpha, 0
+            )
+
     # Benchmark LIF
     torch.cuda.synchronize()
     start_time = time.time()
     iterations = 50
     with torch.no_grad():
         for _ in range(iterations):
-            LIFSubFunction.apply(x_seq, v_init, decay_scalar, threshold_scalar, alpha, 0)
+            LIFSubFunction.apply(
+                x_seq, v_init, decay_scalar, threshold_scalar, alpha, 0
+            )
     torch.cuda.synchronize()
     lif_time = (time.time() - start_time) / iterations
     print(f"LIF Time:  {lif_time * 1000:.3f} ms")
@@ -61,12 +67,15 @@ def benchmark_plif():
     start_time = time.time()
     with torch.no_grad():
         for _ in range(iterations):
-            PLIFSubFunction.apply(x_seq, v_init, decay_vector, threshold_vector, alpha, 0)
+            PLIFSubFunction.apply(
+                x_seq, v_init, decay_vector, threshold_vector, alpha, 0
+            )
     torch.cuda.synchronize()
     plif_time = (time.time() - start_time) / iterations
     print(f"PLIF Time: {plif_time * 1000:.3f} ms")
-    
+
     print(f"Overhead: {plif_time / lif_time:.2f}x slower")
 
+
 if __name__ == "__main__":
     benchmark_plif()

docs/conf.py

@@ -1,6 +1,7 @@
 import os
 import sys
 import re
+
 # Ensure Sphinx can find the source code
 sys.path.insert(0, os.path.abspath("../src"))
 
@@ -37,22 +38,22 @@ def get_project_metadata():
 
 meta = get_project_metadata()
 
-project = 'Tether'
-copyright = '2025, Khushiyant'
+project = "Tether"
+copyright = "2025, Khushiyant"
 author = meta.get("author", "Khushiyant")
 release = meta.get("release", "0.1.0")
 
 extensions = [
-    'sphinx.ext.autodoc',      # Core library for html generation from docstrings
-    'sphinx.ext.autosummary',  # Create neat summary tables
-    'sphinx.ext.napoleon',     # Support for NumPy/Google style docstrings
-    'sphinx.ext.viewcode',     # Add links to highlighted source code
-    'myst_parser',             # Support for Markdown files
-    'sphinx_autodoc_typehints' # Show type hints in docs
+    "sphinx.ext.autodoc",  # Core library for html generation from docstrings
+    "sphinx.ext.autosummary",  # Create neat summary tables
+    "sphinx.ext.napoleon",  # Support for NumPy/Google style docstrings
+    "sphinx.ext.viewcode",  # Add links to highlighted source code
+    "myst_parser",  # Support for Markdown files
+    "sphinx_autodoc_typehints",  # Show type hints in docs
 ]
 
-templates_path = ['_templates']
-exclude_patterns = ['_build', 'Thumbs.db', '.DS_Store']
+templates_path = ["_templates"]
+exclude_patterns = ["_build", "Thumbs.db", ".DS_Store"]
 
 html_static_path = ["_static"]
 html_theme = "furo"
@@ -65,5 +66,5 @@ def get_project_metadata():
 }
 
 suppress_warnings = [
-    'myst.xref_missing',  # Suppress MyST cross-reference errors
-]
+    "myst.xref_missing",  # Suppress MyST cross-reference errors
+]

examples/train_cifar.py

@@ -6,28 +6,26 @@
 from torch.utils.data import DataLoader
 from tether.nn import LIF
 from tether.data import SpikingDatasetWrapper, rate_encoding
-from tether.utils.monitor import Monitor #
+from tether.utils.monitor import Monitor  #
+
 
 class SpikingCIFARModel(nn.Module):
     def __init__(self, n_steps=10):
         super().__init__()
         self.n_steps = n_steps
-        
+
         # Define layers
         self.conv_layers = nn.Sequential(
             nn.Conv2d(3, 32, kernel_size=3, padding=1),
             LIF(32 * 32 * 32),
             nn.MaxPool2d(2),
             nn.Conv2d(32, 64, kernel_size=3, padding=1),
             LIF(64 * 16 * 16),
-            nn.MaxPool2d(2)
+            nn.MaxPool2d(2),
         )
-        
+
         self.fc_layers = nn.Sequential(
-            nn.Flatten(),
-            nn.Linear(64 * 8 * 8, 256),
-            LIF(256),
-            nn.Linear(256, 10)
+            nn.Flatten(), nn.Linear(64 * 8 * 8, 256), LIF(256), nn.Linear(256, 10)
         )
 
     def forward(self, x):
@@ -40,6 +38,7 @@ def forward(self, x):
             outputs.append(out)
         return torch.stack(outputs).mean(0)
 
+
 def evaluate(model, loader, device):
     model.eval()
     correct = 0
@@ -52,7 +51,8 @@ def evaluate(model, loader, device):
             pred = output.argmax(dim=1)
             correct += (pred == target).sum().item()
             total += target.size(0)
-    return 100. * correct / total
+    return 100.0 * correct / total
+
 
 def main():
     # --- Configuration ---
@@ -63,25 +63,36 @@ def main():
     lr = 1e-3
 
     # --- Data Preparation ---
-    transform = transforms.Compose([
-        transforms.ToTensor(),
-        transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
-    ])
-
-    train_raw = datasets.CIFAR10(root="./data", train=True, download=True, transform=transform)
-    test_raw = datasets.CIFAR10(root="./data", train=False, download=True, transform=transform)
-
-    train_ds = SpikingDatasetWrapper(train_raw, encode_fn=lambda x: rate_encoding(x, n_steps=n_steps))
-    test_ds = SpikingDatasetWrapper(test_raw, encode_fn=lambda x: rate_encoding(x, n_steps=n_steps))
-
-    train_loader = DataLoader(train_ds, batch_size=batch_size, shuffle=True, num_workers=4)
-    test_loader = DataLoader(test_ds, batch_size=batch_size, shuffle=False, num_workers=4)
+    transform = transforms.Compose(
+        [transforms.ToTensor(), transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))]
+    )
+
+    train_raw = datasets.CIFAR10(
+        root="./data", train=True, download=True, transform=transform
+    )
+    test_raw = datasets.CIFAR10(
+        root="./data", train=False, download=True, transform=transform
+    )
+
+    train_ds = SpikingDatasetWrapper(
+        train_raw, encode_fn=lambda x: rate_encoding(x, n_steps=n_steps)
+    )
+    test_ds = SpikingDatasetWrapper(
+        test_raw, encode_fn=lambda x: rate_encoding(x, n_steps=n_steps)
+    )
+
+    train_loader = DataLoader(
+        train_ds, batch_size=batch_size, shuffle=True, num_workers=4
+    )
+    test_loader = DataLoader(
+        test_ds, batch_size=batch_size, shuffle=False, num_workers=4
+    )
 
     # --- Model & Monitoring ---
     model = SpikingCIFARModel(n_steps=n_steps).to(device)
     optimizer = optim.AdamW(model.parameters(), lr=lr)
     criterion = nn.CrossEntropyLoss()
-    
+
     # Initialize the Tether Monitor
     monitor = Monitor(model)
 
@@ -115,19 +126,26 @@ def main():
         # --- End of Epoch Monitoring ---
         epoch_time = time.time() - start_time
         avg_loss = running_loss / len(train_loader)
-        train_acc = 100. * train_correct / train_total
+        train_acc = 100.0 * train_correct / train_total
         val_acc = evaluate(model, test_loader, device)
 
         # Retrieve firing rates from all LIF layers via Monitor
         firing_rates = monitor.get_firing_rates()
         # Calculate mean firing rate across the entire model
-        mean_fr = sum(firing_rates.values()) / len(firing_rates) if firing_rates else 0.0
+        mean_fr = (
+            sum(firing_rates.values()) / len(firing_rates) if firing_rates else 0.0
+        )
 
         # Print detailed report once per epoch
-        print(f"Epoch {epoch+1}/{epochs} | Time: {epoch_time:.1f}s")
-        print(f"  > Loss: {avg_loss:.4f} | Train Acc: {train_acc:.2f}% | Val Acc: {val_acc:.2f}%")
-        print(f"  > Mean Firing Rate: {mean_fr:.4f} (Sparsity: {(1-mean_fr)*100:.1f}%)")
+        print(f"Epoch {epoch + 1}/{epochs} | Time: {epoch_time:.1f}s")
+        print(
+            f"  > Loss: {avg_loss:.4f} | Train Acc: {train_acc:.2f}% | Val Acc: {val_acc:.2f}%"
+        )
+        print(
+            f"  > Mean Firing Rate: {mean_fr:.4f} (Sparsity: {(1 - mean_fr) * 100:.1f}%)"
+        )
         print("-" * 60)
 
+
 if __name__ == "__main__":
-    main()
+    main()