updates

vpuri3 · May 13, 2024 · 907a038 · 907a038
1 parent 0232e9e
commit 907a038
Show file tree

Hide file tree

Showing 3 changed files with 13 additions and 21 deletions.
diff --git a/README.md b/README.md
@@ -65,12 +65,12 @@ pK, stK = Lux.setup(rng, kan) |> device
   34.360 μs (175 allocations: 4.78 KiB)
   155.781 μs (565 allocations: 17.50 KiB)
 ```
-With `use_base_activcation = false`, the performance of KAN effectively doubles
+With `use_base_act = false`, the performance of KAN effectively doubles
 ```julia
 kan = Chain(
-    KDense(1, 8, 15; use_base_activation = false),
-    KDense(8, 8, 15; use_base_activation = false),
-    KDense(8, 1, 15; use_base_activation = false),
+    KDense(1, 8, 15; use_base_act = false),
+    KDense(8, 8, 15; use_base_act = false),
+    KDense(8, 1, 15; use_base_act = false),
 )
 p, st = Lux.setup(rng, kan) |> device
 @btime CUDA.@sync $mlp($x, $p, $st)

diff --git a/examples/eg1.jl b/examples/eg1.jl
@@ -27,9 +27,9 @@ function main()
     )
 
     kan = Chain(
-        KDense(1, 8, 15; use_base_act = false),
-        KDense(8, 8, 15; use_base_act = false),
-        KDense(8, 1, 15; use_base_act = false),
+        KDense(1, 8, 15; use_base_act = true),
+        KDense(8, 8, 15; use_base_act = true),
+        KDense(8, 1, 15; use_base_act = true),
     )
 
     # display(mlp)

diff --git a/src/type.jl b/src/type.jl
@@ -73,19 +73,7 @@ function LuxCore.parameterlength(
     len
 end
 
-function (l::KDense{false})(x::AbstractVecOrMat, p, st)
-    K = size(x, 2)                # [I, K]
-    x_resh = reshape(x, 1, :)     # [1, I * K]
-    x_norm = l.normalizer(x_resh) # ∈ [-1, 1]
-
-    basis = rbf.(x_norm, st.grid, l.denominator)      # [G, I * K]
-    basis = reshape(basis, l.grid_len * l.in_dims, K) # [G * I, K]
-    y = p.W1 * basis                                  # [O, K]
-
-    y, st
-end 
-
-function (l::KDense{true})(x::AbstractVecOrMat, p, st)
+function (l::KDense{use_base_act})(x::AbstractVecOrMat, p, st) where{use_base_act}
     K = size(x, 2)                # [I, K]
     x_resh = reshape(x, 1, :)     # [1, I * K]
     x_norm = l.normalizer(x_resh) # ∈ [-1, 1]
@@ -94,7 +82,11 @@ function (l::KDense{true})(x::AbstractVecOrMat, p, st)
     basis = reshape(basis, l.grid_len * l.in_dims, K) # [G * I, K]
     spline = p.W1 * basis                             # [O, K]
 
-    y = spline + p.W2 * l.base_act.(x)
+    y = if use_base_act
+        spline + p.W2 * l.base_act.(x)
+    else
+        spline
+    end
 
     y, st
 end