### Benchmarking KAN and MLP Architectures in Julia

Source: https://github.com/vpuri3/kolmogorovarnold.jl/blob/master/README.md

This comprehensive snippet benchmarks the performance of various Kolmogorov-Arnold Network (KAN) implementations (KDense, CDense, FDense) against a standard Multi-Layer Perceptron (MLP) using Lux.jl and LuxCUDA. It defines each network, initializes their parameters on a GPU, and then measures both forward and backward pass execution times using `BenchmarkTools.@btime` and `Zygote.gradient`.

```Julia
using Lux, KolmogorovArnold
using LuxCUDA

CUDA.allowscalar(false)
device = Lux.gpu_device()

rng = Random.default_rng()
Random.seed!(rng, 0)

x  = rand32(rng, 1, 1000) |> device
x₀ = rand32(rng, 1000, 1) |> device
# define MLP, KANs

mlp = Chain(
    Dense(1, 128, tanh),
    Dense(128, 128, tanh),
    Dense(128, 1),
) # 16_897 parameters plus 0 states.

basis_func = rbf # rbf, rswaf, iqf (radial basis funcs, reflection switch activation funcs, inverse quadratic funcs)
normalizer = softsign # sigmoid(_fast), tanh(_fast), softsign

kan1 = Chain(
    KDense( 1, 40, 10; use_base_act = true, basis_func, normalizer),
    KDense(40, 40, 10; use_base_act = true, basis_func, normalizer),
    KDense(40,  1, 10; use_base_act = true, basis_func, normalizer),
) # 18_490 parameters plus 30 states.

kan2 = Chain(
    KDense( 1, 40, 10; use_base_act = false, basis_func, normalizer),
    KDense(40, 40, 10; use_base_act = false, basis_func, normalizer),
    KDense(40,  1, 10; use_base_act = false, basis_func, normalizer),
) # 16_800 parameters plus 30 states.

kan3 = Chain(
    CDense( 1, 40, G),
    CDense(40, 40, G),
    CDense(40,  1, G),
) # 18_561 parameters plus 0 states.

kan4 = Chain(
    FDense( 1, 30, G),
    FDense(30, 30, G),
    FDense(30,  1, G),
) # 19_261 parameters plus 0 states.

# set up experiment
pM, stM = Lux.setup(rng, mlp)
pK1, stK1 = Lux.setup(rng, kan1)
pK2, stK2 = Lux.setup(rng, kan2)
pK3, stK3 = Lux.setup(rng, kan3)
pK4, stK4 = Lux.setup(rng, kan4)


pM  = ComponentArray(pM) |> device
pK1 = ComponentArray(pK1) |> device
pK2 = ComponentArray(pK2) |> device
pK3 = ComponentArray(pK3) |> device
pK4 = ComponentArray(pK4) |> device

stM, stK1, stK2, stK3, stK4 = device(stM), device(stK1), device(stK2), device(stK4), device(stK4)

# Forward pass

@btime CUDA.@sync $mlp($x, $pM, $stM)     # 31.611 μs (248 allocations: 5.45 KiB)
@btime CUDA.@sync $kan1($x, $pK1, $stK1)  # 125.790 μs (1034 allocations: 21.97 KiB)
@btime CUDA.@sync $kan2($x, $pK2, $stK2)  # 87.585 μs (1335 allocations: 13.95 KiB)
@btime CUDA.@sync $kan3($x', $pK3, $stK3) # 210.785 μs (1335 allocations: 31.03 KiB)
@btime CUDA.@sync $kan4($x', $pK4, $stK4) # 2.392 ms (1642 allocations: 34.56 KiB)


# Backward pass

f_mlp(p)  = mlp(x, p, stM)[1] |> sum
f_kan1(p) = kan1(x, p, stK1)[1] |> sum
f_kan2(p) = kan2(x, p, stK2)[1] |> sum
f_kan3(p) = kan3(x₀, p, stK3)[1] |> sum
f_kan4(p) = kan4(x₀, p, stK4)[1] |> sum


@btime CUDA.@sync Zygote.gradient($f_mlp, $pM)   # 268.074 μs (1971 allocations: 57.03 KiB)
@btime CUDA.@sync Zygote.gradient($f_kan1, $pK1) # 831.888 μs (5015 allocations: 123.25 KiB)
@btime CUDA.@sync Zygote.gradient($f_kan2, $pK2) # 658.578 μs (3314 allocations: 87.16 KiB)
@btime CUDA.@sync Zygote.gradient($f_kan3, $pK3) # 1.647 ms (7138 allocations: 180.45 KiB)
@btime CUDA.@sync Zygote.gradient($f_kan4, $pK4) # 7.028 ms (8745 allocations: 199.42 KiB)
```

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

### Initializing and Forward Pass with KDense Layer in Julia

Source: https://github.com/vpuri3/kolmogorovarnold.jl/blob/master/README.md

This snippet demonstrates the basic initialization and forward pass of a `KDense` layer, a core component of Kolmogorov-Arnold Networks, within the Lux.jl framework. It sets up a random number generator, defines layer dimensions, initializes the layer's parameters and states, and then performs a forward pass with random input data.

```Julia
using Random, KolmogorovArnold
rng = Random.default_rng()

in_dim, out_dim, grid_len = 4, 4, 8
layer = KDense(in_dim, out_dim, grid_len)
p, st = Lux.setup(rng, layer)

x = rand32(rng, in_dim, 10)
y = layer(x, p, st)
```

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

### Benchmarking Custom Gradients for Activation Functions in Julia

Source: https://github.com/vpuri3/kolmogorovarnold.jl/blob/master/README.md

This snippet defines and benchmarks the forward and backward passes for RBF, RSWAF, and IQF activation functions using custom gradients. It initializes input arrays `x` and `z` on the device, calculates function sums, and then measures performance with `@btime` for both forward and `Zygote.gradient` for backward passes, highlighting the speedup achieved by custom implementations.

```Julia
N, G = 5000, 10

x = LinRange(-1, 1, N) |> Array |> device
z = LinRange(-1, 1, G) |> Array |> device
d = 2 / (G-1)

f_rbf(z)   = rbf(  x, z', d) |> sum
f_rswaf(z) = rswaf(x, z', d) |> sum
f_iqf(z)   = iqf(  x, z', d) |> sum

# forward pass
@btime CUDA.@sync $f_rbf($z)   # 55.566 μs (294 allocations: 7.78 KiB)
@btime CUDA.@sync $f_rswaf($z) # 57.112 μs (294 allocations: 7.78 KiB)
@btime CUDA.@sync $f_iqf($z)   # 55.368 μs (294 allocations: 7.78 KiB)

# backward pass
@btime CUDA.@sync Zygote.gradient($f_rbf  , $z) # 188.456 μs (1045 allocations: 27.62 KiB)
@btime CUDA.@sync Zygote.gradient($f_rswaf, $z) # 212.419 μs (1071 allocations: 28.30 KiB)
@btime CUDA.@sync Zygote.gradient($f_iqf  , $z) # 201.568 μs (1045 allocations: 27.62 KiB)

# without custom gradients
# RBF  : 250.313 μs (1393 allocations: 35.95 KiB)
# RSWAF: 282.864 μs (1389 allocations: 36.62 KiB)
# IQF  : 333.843 μs (1628 allocations: 42.70 KiB)
```

=== COMPLETE CONTENT === This response contains all available snippets from this library. No additional content exists. Do not make further requests.