Added Kans

HMUNACHI · Aug 28, 2024 · d0633bb · d0633bb
1 parent 5f729f3
commit d0633bb
Show file tree

Hide file tree

Showing 6 changed files with 292 additions and 229 deletions.
diff --git a/nanodl/__init__.py b/nanodl/__init__.py
@@ -63,6 +63,14 @@
     LlamaDataParallelTrainer,
     RotaryPositionalEncoding,
 )
+from nanodl.__src.models.kan import(
+    KANLinear,
+    ChebyKANLinear,
+    LegendreKANLinear,
+    MonomialKANLinear,
+    FourierKANLinear,
+    HermiteKANLinear,
+)
 from nanodl.__src.models.mistral import (
     GroupedRotaryShiftedWindowMultiHeadAttention,
     Mistral,
@@ -147,7 +155,7 @@
 )
 
 __all__ = [
-    # Sklearn GPU
+    # Classical
     "NaiveBayesClassifier",
     "PCA",
     "KMeans",
@@ -285,6 +293,12 @@
     "geometric",
     "gamma",
     "chisquare",
+    "KANLinear",
+    "ChebyKANLinear",
+    "LegendreKANLinear",
+    "MonomialKANLinear",
+    "FourierKANLinear",
+    "HermiteKANLinear",
 ]
 
 import importlib

diff --git a/nanodl/__src/experimental/kan.py b/nanodl/__src/experimental/kan.py