Implemented FastKAN in TinyGrad

[mdaiter]

Hey hey!

Nice work with FastKAN. I got this running on Metal in TinyGrad, if it's of any interest. Really cool results, and very fast.

The only thing that's still a bit wonky / broken is plotting. I can take it out, if it helps merge it.

Best

[mdaiter]

Also:

from tinygrad import TinyJit
with Context(BEAM=2):
    fklayer = FastKANLayer(100, 100)
    x = Tensor.randn(8, 100)
    rbf = RadialBasisFunction()
    def fklayer_step():
        fklayer(x, use_layernorm=False)
    def fklayer_sum_backwards_test():
        fklayer(x).sum().backward()
    def rbf_test():
        rbf(x)
    %timeit -r10 -n1000 TinyJit(fklayer_step)
    %timeit -r10 -n1000 TinyJit(fklayer_sum_backwards_test)
    %timeit -r10 -n1000 TinyJit(rbf_test)

Results:

345 ns ± 157 ns per loop (mean ± std. dev. of 10 runs, 1,000 loops each)
304 ns ± 24.4 ns per loop (mean ± std. dev. of 10 runs, 1,000 loops each)
334 ns ± 43 ns per loop (mean ± std. dev. of 10 runs, 1,000 loops each)

^ran on Macbook Air with a M2 processor, stupidly fast :)