KANNs

%23%20%2F%2F%2F%20script%0A%23%20requires-python%20%3D%20%22%3E%3D3.13%22%0A%23%20dependencies%20%3D%20%5B%22marimo%3E%3D0.23.5%22%2C%20%22matplotlib%3D%3D3.10.8%22%2C%20%22numpy%3D%3D2.4.3%22%2C%20%22sympy%3D%3D1.14.0%22%2C%20%22torch%3D%3D2.9.1%22%5D%0A%23%20%0A%23%20%5Btool.marimo.runtime%5D%0A%23%20cache_cells%20%3D%20true%0A%23%20pin_modules%20%3D%20false%0A%23%20%0A%23%20%5Btool.marimo.export%5D%0A%23%20lock_kind%20%3D%20%22observed%22%0A%23%20%2F%2F%2F%0A%0Aimport%20marimo%0A%0A__generated_with%20%3D%20%220.23.13%22%0Aapp%20%3D%20marimo.App(width%3D%22medium%22%2C%20app_title%3D%22KANNs%22)%0A%0Awith%20app.setup%3A%0A%20%20%20%20import%20matplotlib.pyplot%20as%20plt%0A%20%20%20%20import%20numpy%20as%20np%0A%0A%20%20%20%20import%20marimo%20as%20mo%0A%20%20%20%20import%20torch%0A%20%20%20%20import%20torch.nn%20as%20nn%0A%20%20%20%20import%20torch.optim%20as%20optim%0A%20%20%20%20from%20typing%20import%20Union%0A%20%20%20%20import%20math%0A%0A%20%20%20%20%23%20Discrete%20UI%20state%20space.%20The%20UI%20elements%20below%20expose%20*indices*%20into%0A%20%20%20%20%23%20these%20lists%20(not%20the%20values%20themselves)%2C%20so%20the%20full%20interactive%20state%0A%20%20%20%20%23%20space%20is%20a%20finite%20product%20of%20index%20ranges%20that%20a%20precompute%20sweep%20can%0A%20%20%20%20%23%20enumerate%20and%20cache%20exhaustively.%0A%20%20%20%20FN_OPTIONS%20%3D%20%5B%0A%20%20%20%20%20%20%20%20(%22x%5E2%22%2C%20lambda%20x%3A%20x**2)%2C%0A%20%20%20%20%20%20%20%20(%22sin(x%5E2)%20*%20x%22%2C%20lambda%20x%3A%20torch.sin(x**2)%20*%20x)%2C%0A%20%20%20%20%5D%0A%20%20%20%20%23%20Labels%20alone%2C%20split%20out%20from%20the%20callables.%20Cells%20that%20only%20need%20to%0A%20%20%20%20%23%20*display*%20the%20option%20names%20(the%20UI%20dropdown%2C%20sweep%20bounds)%20reference%20this%0A%20%20%20%20%23%20serializable%20list%20instead%20of%20FN_OPTIONS%2C%20so%20a%20cached%20restore%20hands%20them%20a%0A%20%20%20%20%23%20real%20list%20rather%20than%20an%20unpicklable-lambda%20stub%20%E2%80%94%20the%20lambdas%20stay%20in%0A%20%20%20%20%23%20FN_OPTIONS%2C%20touched%20only%20inside%20the%20cached%20sweep%20body.%0A%20%20%20%20FN_LABELS%20%3D%20%5Blabel%20for%20label%2C%20_%20in%20FN_OPTIONS%5D%0A%20%20%20%20LAYER_OPTIONS%20%3D%20%5B%5B3%5D%2C%20%5B2%5D%2C%20%5B(2%2C%201)%5D%2C%20%5B(4%2C%202)%5D%5D%0A%20%20%20%20%23%20Slider%20indices%200..NOISE_STEPS-1%20map%20linearly%20onto%20noise%20in%20%5B0%2C%201%5D.%0A%20%20%20%20NOISE_STEPS%20%3D%2011%0A%0A%20%20%20%20def%20noise_at(index%3A%20int)%20-%3E%20float%3A%0A%20%20%20%20%20%20%20%20return%20index%20%2F%20(NOISE_STEPS%20-%201)%0A%0A%20%20%20%20def%20is_wasm()%20-%3E%20bool%3A%0A%20%20%20%20%20%20%20%20%23%20True%20inside%20the%20browser%20(Pyodide)%20kernel%2C%20where%20torch%20is%20absent%20and%0A%20%20%20%20%20%20%20%20%23%20every%20result%20is%20served%20from%20the%20bundled%20cache%20instead%20of%20computed.%0A%20%20%20%20%20%20%20%20import%20sys%0A%0A%20%20%20%20%20%20%20%20return%20sys.platform%20%3D%3D%20%22emscripten%22%0A%0A%0A%40app.cell(hide_code%3DTrue)%0Adef%20_()%3A%0A%20%20%20%20mo.md(r%22%22%22%0A%20%20%20%20%23%20Kolomogrov-Arnold%20Networks%20(KANs)%0A%0A%20%20%20%20The%20inference%20step%20of%20a%20classic%20neural%20network%20is%20a%20series%20of%20differentiable%0A%20%20%20%20composed%20functions%2C%20e.g.%0A%0A%20%20%20%20%24%24%5Cbar%20Y_p%20%3D%20W_0%20f_0(W_1%2C%20f_1(W_2%2C%20f_2...%20f_i(W_i%2C%20%5Cbar%20x)...))%24%24%0A%0A%20%20%20%20Kolonogrov-Arnold%20Neural%20Networks%20(KANNs)%2C%20subvert%20this%20classic%20composition%20by%0A%20%20%20%20using%20weights%20to%20produce%20a%20function%2C%20and%20not%20direct%20an%20output.%20That%20is%0A%0A%20%20%20%20%24%24%5Cbar%20Y_p%20%3D%20(g_0(W_0))%20%5Ccirc%20%5Cleft%5B%5Csum%20(g_1(W_1))%20%5Ccirc%20%5Cleft%5B...%20%5Csum%20(g_i(W_i))%20%5Ccirc%20%5B%5Cbar%20x%5D%5Cright%5D%5Cright%5D%20%24%24%0A%0A%20%20%20%20so%20while%20%24f%3A%20%5Cleft(x%20%5Cin%20%5Cmathbb%7BR%5Em%7D%2C%20W_f%20%5Cin%20%5Cmathbb%7BR%5E%7Bm%20%5Ctimes%20n%7D%7D%5Cright)%20%5Crightarrow%20y%20%5Cin%20%5Cmathbb%7BR%5En%7D%24%20is%20typically%20a%20predefined%20function%20or%20%22activation%22%20function%20(such%20as%20ReLU%20or%20tanh)%20and%20%24m%2Cn%24%20are%20the%20dimensions%20of%20the%20input%20and%20output%20respectively.%0A%20%20%20%20In%20KANs%20%24g_n%20%3A%20%5Cleft(W_g%20%5Cin%20%5Cmathbb%7BR%5E%7Bk%20%5Ctimes%20m%20%5Ctimes%20n%7D%7D%20%5Cright)%20%5Crightarrow%20(x%5Cin%20%5Cmathbb%7BR%5Em%7D%5Crightarrow%20y%20%5Cin%20%5Cmathbb%7BR%5En%7D)%24%2C%20where%20%24g%24%20is%20typically%20a%20spline%20function%2C%20and%20%24k%24%20is%20the%20number%20of%20evaluation%20points.%0A%20%20%20%20Here%2C%20I%20denote%20composition%20as%20%24%5Ccirc%24%2C%20that%20is%20%24F(G(x))%20%3D%20(F)%20%5Ccirc%20%5BG(x)%5D%24%2C%20in%20order%20to%0A%20%20%20%20express%20that%20%24g_n%24%20is%20a%20function%20that%20produces%20a%20function%20and%20not%20a%20direct%0A%20%20%20%20output.%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell(hide_code%3DTrue)%0Adef%20_()%3A%0A%20%20%20%20mo.md(r%22%22%22%0A%20%20%20%20%23%23%20Baseline%20comparison%0A%0A%20%20%20%20%23%23%23%20Case%201%0A%0A%20%20%20%20For%20a%20baseline%20comparison%20between%20classic%20NNs%20and%20KANs%2C%20the%20number%20of%20free%20weights%20should%20be%20the%20same.%0A%20%20%20%20We%20notice%20that%20KANs%20have%20the%20extra%20hyperparameter%20of%20%24k%24%2C%20which%20is%20the%20number%20of%20evaluation%20points.%0A%20%20%20%20Letting%20%24k%3Dm%24%2C%20allows%20the%20weights%20of%20%24g%24%2C%20%24W_g%5Cin%20%5Cmathbb%7BR%5E%7Bm%20%5Ctimes%20m%20%5Ctimes%20n%7D%7D%24%2C%20be%20considered%20as%20a%20matrix%20with%20dimensions%20%24(m%2C%20n)%24%20repeated%20m%20times%2C%20and%20thus%20having%20%24mn%24%20free%20variables%20as%20in%20the%20classical%20case.%0A%0A%20%20%20%20%23%23%23%20Case%202%0A%0A%20%20%20%20Conversely%2C%20for%20a%20fixed%20%24k%24%20let%20%24W_f%5Cin%5Cbegin%7Bcases%7D%0A%20%20%20%20i%7C2%2C%20%26%20%5Cmathbb%7BR%5E%7Bm%5Ctimes%20n%20k%7D%7D%5C%5C%0A%20%20%20%20otherwise%2C%20%26%20%5Cmathbb%7BR%5E%7Bk%20m%5Ctimes%20n%7D%7D%0A%20%20%20%20%5Cend%7Bcases%7D%24%20where%20%24i%24%20is%20layer%20index%3B%20or%2C%20if%20%24k%24%20is%20a%20perfect%20square%2C%20let%20%24W_f%5Cin%5Cbegin%7Bcases%7D%0A%20%20%20%20i%7C2%2C%20%26%20%5Cmathbb%7BR%5E%7Bm%5Ctimes%20%5Csqrt%20k%7D%7D%5C%5C%0A%20%20%20%20otherwise%2C%20%26%20%5Cmathbb%7BR%5E%7B%5Csqrt%20k%5Ctimes%20n%7D%7D%0A%20%20%20%20%5Cend%7Bcases%7D%24%2C%20for%20twice%20the%20number%20of%20layers%3B%20allowing%20%24kmn%24%20free%20variables.%20This%20write-up%20ignores%20this%20latter%20case%2C%20as%20potentially%20contracting%20the%20layer%20dimension%20can%20have%20unintended%20consequences.%0A%0A%20%20%20%20%23%23%23%20Evaluation%0A%0A%20%20%20%20To%20explore%20the%20relationship%20between%20classical%20neural%20networks%20and%20KANs%2C%20we%20examine%20a%20couple%20cases%20below.%20In%20all%20the%20target%20functions%20listed%20noise%20is%20pulled%20from%20a%20normal%20distribution%3B%20and%20layers%20take%20either%20the%20form%20%24%5Bm%2C...%5D%24%20or%20%24%5B(k%2C%20m)%2C...%5D%5D%24%20from%20equivalence%20case%201%20and%202%20respectively.%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell(hide_code%3DTrue)%0Adef%20_()%3A%0A%20%20%20%20%23%20Each%20control's%20%60.value%60%20is%20an%20*index*%20into%20the%20setup%20option%20lists.%0A%20%20%20%20fn%20%3D%20mo.ui.dropdown(%0A%20%20%20%20%20%20%20%20options%3D%7Blabel%3A%20i%20for%20i%2C%20label%20in%20enumerate(FN_LABELS)%7D%2C%0A%20%20%20%20%20%20%20%20value%3D%22sin(x%5E2)%20*%20x%22%2C%0A%20%20%20%20%20%20%20%20label%3D%22a%20sample%20function%22%2C%0A%20%20%20%20)%0A%20%20%20%20layers%20%3D%20mo.ui.dropdown(%0A%20%20%20%20%20%20%20%20options%3D%7Bstr(l)%3A%20i%20for%20i%2C%20l%20in%20enumerate(LAYER_OPTIONS)%7D%2C%0A%20%20%20%20%20%20%20%20value%3Dstr(LAYER_OPTIONS%5B0%5D)%2C%0A%20%20%20%20%20%20%20%20label%3D%22Layers%20for%20method%22%2C%0A%20%20%20%20)%0A%20%20%20%20noise%20%3D%20mo.ui.slider(%0A%20%20%20%20%20%20%20%20start%3D0%2C%0A%20%20%20%20%20%20%20%20stop%3DNOISE_STEPS%20-%201%2C%0A%20%20%20%20%20%20%20%20step%3D1%2C%0A%20%20%20%20%20%20%20%20label%3D%22noise%22%2C%0A%20%20%20%20%20%20%20%20value%3Dround(0.2%20*%20(NOISE_STEPS%20-%201))%2C%0A%20%20%20%20)%0A%20%20%20%20mo.vstack(%5Bfn%2C%20layers%2C%20noise%5D)%0A%20%20%20%20return%20fn%2C%20layers%2C%20noise%0A%0A%0A%40app.cell%0Adef%20plot(result)%3A%0A%20%20%20%20plt.title(f%22%24%7Bresult%5B'fn_label'%5D%7D%24%22)%0A%20%20%20%20plt.scatter(%0A%20%20%20%20%20%20%20%20result%5B%22x_np%22%5D%2C%0A%20%20%20%20%20%20%20%20result%5B%22y_np%22%5D%2C%0A%20%20%20%20%20%20%20%20label%3D%22raw%20data%22%2C%0A%20%20%20%20%20%20%20%20s%3D5%2C%0A%20%20%20%20%20%20%20%20facecolors%3D%22none%22%2C%0A%20%20%20%20%20%20%20%20edgecolors%3D%22grey%22%2C%0A%20%20%20%20)%0A%20%20%20%20plt.scatter(%0A%20%20%20%20%20%20%20%20result%5B%22x_np%22%5D%2C%0A%20%20%20%20%20%20%20%20result%5B%22classic_pred%22%5D%2C%0A%20%20%20%20%20%20%20%20label%3Df%22Classic%20NN%20(weights%20%7Bresult%5B'layers_label'%5D%7D)%22%2C%0A%20%20%20%20%20%20%20%20s%3D7%2C%0A%20%20%20%20)%0A%20%20%20%20plt.scatter(%0A%20%20%20%20%20%20%20%20result%5B%22x_np%22%5D%2C%0A%20%20%20%20%20%20%20%20result%5B%22kan_pred%22%5D%2C%0A%20%20%20%20%20%20%20%20label%3Df%22KAN%20(weights%20%7Bresult%5B'layers_label'%5D%7D)%22%2C%0A%20%20%20%20%20%20%20%20s%3D7%2C%0A%20%20%20%20)%0A%20%20%20%20plt.xlabel(%22x%22)%0A%20%20%20%20plt.ylabel(%22y%22)%0A%20%20%20%20plt.legend()%0A%20%20%20%20plt.gca()%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(result)%3A%0A%20%20%20%20import%20sympy%0A%20%20%20%20from%20sympy.abc%20import%20X%0A%0A%0A%20%20%20%20%23%20%60weight%60%20is%20a%20numpy%20array%20(the%20KAN%20parameters%20were%20detached%20to%20numpy%20in%0A%20%20%20%20%23%20%60kann_permutation%60)%2C%20so%20this%20cell%20renders%20symbolically%20without%20torch.%0A%20%20%20%20def%20w_to_sympy_2(X%2C%20weight)%3A%0A%20%20%20%20%20%20%20%20bias%20%3D%20np.array(weight%5B0%2C%20%3A%5D)%0A%20%20%20%20%20%20%20%20weights%20%3D%20np.array(weight%5B1%3A%2C%20%3A%5D)%0A%20%20%20%20%20%20%20%20if%20weight.shape%5B0%5D%20%3D%3D%202%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20span%20%3D%20np.array(%5B0%5D)%0A%20%20%20%20%20%20%20%20else%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20span%20%3D%20np.linspace(-1%2C%201%2C%20weight.shape%5B0%5D%20-%201)%0A%20%20%20%20%20%20%20%20X%20%3D%20sympy.simplify(%0A%20%20%20%20%20%20%20%20%20%20%20%20(span%20-%20X)%20**%20np.arange(1%2C%20weight.shape%5B0%5D)%20%40%20weights%20%2B%20bias%0A%20%20%20%20%20%20%20%20)%0A%20%20%20%20%20%20%20%20return%20np.reshape(X%2C%20(-1%2C%20weight.shape%5B1%5D))%0A%0A%0A%20%20%20%20def%20w_to_sympy(X%2C%20weight)%3A%0A%20%20%20%20%20%20%20%20if%20weight.ndim%20%3D%3D%202%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20return%20w_to_sympy_2(X%2C%20weight)%0A%20%20%20%20%20%20%20%20k%20%3D%20weight.shape%5B0%5D%20-%201%0A%20%20%20%20%20%20%20%20span%20%3D%20np.linspace(-1%2C%201%2C%20k)%0A%20%20%20%20%20%20%20%20batch_size%20%3D%20len(X)%0A%20%20%20%20%20%20%20%20dims%20%3D%20len(X%5B0%5D)%0A%20%20%20%20%20%20%20%20assert%20dims%20%3D%3D%20weight.shape%5B1%5D%0A%20%20%20%20%20%20%20%20x%20%3D%20np.concatenate(%0A%20%20%20%20%20%20%20%20%20%20%20%20%5B%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20np.ones((batch_size%2C%20dims%2C%201))%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20(span%20-%20np.reshape(X%2C%20(batch_size%2C%20dims%2C%201)))%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20**%20np.arange(1%2C%20k%20%2B%201)%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%5D%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20axis%3D2%2C%0A%20%20%20%20%20%20%20%20)%0A%20%20%20%20%20%20%20%20X%20%3D%20np.einsum(%22bmk%2Ckmn-%3Ebn%22%2C%20x%2C%20weight)%0A%20%20%20%20%20%20%20%20return%20sympy.simplify(np.reshape(X%2C%20(-1%2C%20weight.shape%5B1%5D)))%0A%0A%0A%20%20%20%20from%20sympy%20import%20latex%0A%0A%20%20%20%20bullet%20%3D%20%22%24%5Cn%20-%20%24%22%0A%20%20%20%20mo.md(%0A%20%20%20%20%20%20%20%20%22%23%23%23%23%20KANNs%20Results%3A%5Cn%22%0A%20%20%20%20%20%20%20%20%2B%20%22%5Cn%5Cn%22.join(%0A%20%20%20%20%20%20%20%20%20%20%20%20%5B%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20name%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%2B%20%22%3A%5Cn%5Cn%20-%20%24%22%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%2B%20f%22%7Bbullet.join(latex(eq%5B0%5D)%20for%20eq%20in%20w_to_sympy(%5B%5BX%5D%5D%2C%20weight))%7D%24%22%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20for%20name%2C%20weight%20in%20result%5B%22kan_weights%22%5D%0A%20%20%20%20%20%20%20%20%20%20%20%20%5D%0A%20%20%20%20%20%20%20%20)%0A%20%20%20%20)%0A%20%20%20%20return%0A%0A%0A%40app.cell(hide_code%3DTrue)%0Adef%20_()%3A%0A%20%20%20%20mo.md(r%22%22%22%0A%20%20%20%20%23%23%20Some%20observations%20on%20implementation%0A%0A%20%20%20%20KANNs%20seem%20very%20sensitive%20to%20over-fitting%20with%20an%20increased%20number%20of%20parameters.%0A%20%20%20%20Matching%20the%20layers%20with%20the%20expected%20equation%20form%20seems%20to%20perform%20well%2C%20whereas%20just%20providing%20a%20general%20depth%20without%20factoring%20in%20the%20equation%20form%20leads%20to%20strange%20fits.%0A%20%20%20%20This%20makes%20sense%2C%20but%20it%20may%20require%20more%20domain%20knowledge%20of%20the%20problem.%0A%20%20%20%20Anecdotally%2C%20KANNs%20are%20very%20capable%20of%20matching%20systems%2C%20which%20can%20be%20easily%20approximated%20with%20polynomials%20with%20a%20few%20parameters%20over%20standard%20neural%20networks.%0A%0A%20%20%20%20%23%23%20Next%20steps%20for%20simple%20study%0A%0A%20%20%20%20Apply%20KANNs%20to%20PINNs.%20The%20original%20paper%20talks%20about%20this.%20The%20paper%20would%20be%20a%20parametric%20study%20of%20basis%20functions%20against%20equivalent%20classic%20NNs.%0A%0A%20%20%20%20This%20would%20just%20be%20benchmarked%20against%20the%20original%20Kardinakis%20data.%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_()%3A%0A%20%20%20%20def%20global_k(k%2C%20layers)%3A%0A%20%20%20%20%20%20%20%20%23%20If%20a%20consistent%20k%20is%20given%2C%20assume%20that%20is%20the%20global%20k%0A%20%20%20%20%20%20%20%20if%20k%20%3D%3D%20-1%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20for%20layer%20in%20layers%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20if%20isinstance(layer%2C%20tuple)%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20_k%2C%20_%20%3D%20layer%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20if%20k%20%3D%3D%20_k%20or%20k%20%3D%3D%20-1%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20k%20%3D%20_k%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20continue%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20return%20-1%0A%20%20%20%20%20%20%20%20return%20k%0A%0A%0A%20%20%20%20class%20ClassicNN(nn.Module)%3A%0A%20%20%20%20%20%20%20%20def%20__init__(%0A%20%20%20%20%20%20%20%20%20%20%20%20self%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20layers%3A%20list%5Bint%5D%20%3D%20None%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20activation_fn%3Dnn.ELU%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20k%3A%20int%20%3D%20-1%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20output_dim%3A%20int%20%3D%201%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20input_dim%3A%20int%20%3D%201%2C%0A%20%20%20%20%20%20%20%20)%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20super().__init__()%0A%20%20%20%20%20%20%20%20%20%20%20%20if%20layers%20is%20None%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20layers%20%3D%20%5B%5D%0A%0A%20%20%20%20%20%20%20%20%20%20%20%20k%20%3D%20global_k(k%2C%20layers)%0A%0A%20%20%20%20%20%20%20%20%20%20%20%20layers%20%3D%20%5B*layers%2C%20output_dim%5D%0A%20%20%20%20%20%20%20%20%20%20%20%20sequence%20%3D%20%5B%5D%0A%20%20%20%20%20%20%20%20%20%20%20%20prev_dim%20%3D%20input_dim%0A%20%20%20%20%20%20%20%20%20%20%20%20for%20i%2C%20layer%20in%20enumerate(layers)%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20_k%20%3D%20k%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20if%20isinstance(layer%2C%20tuple)%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20_k%2C%20layer%20%3D%20layer%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20dim%20%3D%20k%20*%20layer%20if%20(k%20%3E%200%20and%20i%20%25%202%20%3D%3D%200)%20else%20layer%0A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%23%20If%20k%20doesn't%20match%20the%20expected%20k%2C%20do%20the%20trick%20for%20an%20equal%20number%20of%20weights.%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20if%20_k%20!%3D%20k%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20_sk%20%3D%20math.isqrt(_k)%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20assert%20(%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20_k%20%3D%3D%20_sk**2%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20)%2C%20%22specified%20k%20that%20does%20not%20match%20global%20k%20only%20works%20for%20square%20numbers.%22%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20sequence.extend(%5Bnn.Linear(prev_dim%2C%20_sk)%2C%20activation_fn()%5D)%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20prev_dim%20%3D%20_sk%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20dim%20%3D%20layer%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20sequence.extend(%5Bnn.Linear(prev_dim%2C%20dim)%2C%20activation_fn()%5D)%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20prev_dim%20%3D%20dim%0A%0A%20%20%20%20%20%20%20%20%20%20%20%20self.layers_list%20%3D%20layers%0A%20%20%20%20%20%20%20%20%20%20%20%20self.k%20%3D%20k%0A%20%20%20%20%20%20%20%20%20%20%20%20self.output_dim%20%3D%20output_dim%0A%20%20%20%20%20%20%20%20%20%20%20%20self.model%20%3D%20nn.Sequential(*sequence)%0A%0A%20%20%20%20%20%20%20%20def%20forward(self%2C%20x)%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20x%20%3D%20self.model(x)%0A%20%20%20%20%20%20%20%20%20%20%20%20%23%20If%20we%20output%20on%20an%20extended%20layer%2C%20sum%20back%20to%20the%20correct%20shape.%0A%20%20%20%20%20%20%20%20%20%20%20%20if%20self.k%20%3E%200%20and%20len(self.layers_list)%20%25%202%20%3D%3D%201%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20x%20%3D%20torch.sum(x.reshape(-1%2C%20self.k%2C%20self.output_dim)%2C%20dim%3D1)%0A%20%20%20%20%20%20%20%20%20%20%20%20return%20x%0A%0A%0A%20%20%20%20def%20build_classic_nn_model(%0A%20%20%20%20%20%20%20%20layers%3A%20list%5Bint%5D%20%3D%20None%2C%0A%20%20%20%20%20%20%20%20activation_fn%3Dnn.ELU%2C%0A%20%20%20%20%20%20%20%20seed%3A%20int%20%3D%200%2C%0A%20%20%20%20%20%20%20%20k%3A%20int%20%3D%20-1%2C%0A%20%20%20%20%20%20%20%20output_dim%3A%20int%20%3D%201%2C%0A%20%20%20%20%20%20%20%20input_dim%3A%20int%20%3D%201%2C%0A%20%20%20%20)%3A%0A%20%20%20%20%20%20%20%20torch.manual_seed(seed)%0A%20%20%20%20%20%20%20%20model%20%3D%20ClassicNN(%0A%20%20%20%20%20%20%20%20%20%20%20%20layers%3Dlayers%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20activation_fn%3Dactivation_fn%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20k%3Dk%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20output_dim%3Doutput_dim%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20input_dim%3Dinput_dim%2C%0A%20%20%20%20%20%20%20%20)%0A%20%20%20%20%20%20%20%20return%20model%0A%0A%0A%20%20%20%20class%20KAN(nn.Module)%3A%0A%20%20%20%20%20%20%20%20def%20__init__(%0A%20%20%20%20%20%20%20%20%20%20%20%20self%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20layers%3A%20Union%5Btuple%5BUnion%5Bint%2C%20tuple%5Bint%2C%20int%5D%5D%5D%2C%20list%5D%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20input_dim%3A%20int%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20output_dim%3A%20int%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20seed%3A%20int%20%3D%200%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20k%3A%20int%20%3D%20-1%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20span%3A%20tuple%5Bfloat%2C%20float%5D%20%3D%20(-1.0%2C%201.0)%2C%0A%20%20%20%20%20%20%20%20)%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20super().__init__()%0A%20%20%20%20%20%20%20%20%20%20%20%20torch.manual_seed(seed)%0A%0A%20%20%20%20%20%20%20%20%20%20%20%20self.span%20%3D%20span%0A%20%20%20%20%20%20%20%20%20%20%20%20self.K%20%3D%20k%0A%20%20%20%20%20%20%20%20%20%20%20%20self.input_dim%20%3D%20input_dim%0A%20%20%20%20%20%20%20%20%20%20%20%20self.output_dim%20%3D%20output_dim%0A%0A%20%20%20%20%20%20%20%20%20%20%20%20weights%20%3D%20nn.ParameterList()%0A%20%20%20%20%20%20%20%20%20%20%20%20previous%20%3D%20input_dim%20%2B%201%0A%20%20%20%20%20%20%20%20%20%20%20%20for%20i%2C%20layer%20in%20enumerate(layers)%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20_k%20%3D%20k%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20if%20isinstance(layer%2C%20tuple)%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20_k%2C%20layer%20%3D%20layer%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20dim%20%3D%20(%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20(_k%20%2B%201%2C%20previous%20-%201%2C%20layer)%20if%20_k%20%3E%200%20else%20(previous%2C%20layer)%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20)%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20weights.append(nn.Parameter(torch.empty(dim).uniform_(-1%2C%201)))%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20previous%20%3D%20layer%20%2B%201%0A%0A%20%20%20%20%20%20%20%20%20%20%20%20dim%20%3D%20(%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20(k%20%2B%201%2C%20previous%20-%201%2C%20output_dim)%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20if%20k%20%3E%200%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20else%20(previous%2C%20output_dim)%0A%20%20%20%20%20%20%20%20%20%20%20%20)%0A%20%20%20%20%20%20%20%20%20%20%20%20weights.append(nn.Parameter(torch.empty(dim).uniform_(-1%2C%201)))%0A%20%20%20%20%20%20%20%20%20%20%20%20self.weights%20%3D%20weights%0A%0A%20%20%20%20%20%20%20%20def%20_classic_equivalent(self%2C%20x%2C%20weight)%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20if%20weight.shape%5B0%5D%20%3D%3D%202%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20span%20%3D%20torch.tensor(%5B0.0%5D)%0A%20%20%20%20%20%20%20%20%20%20%20%20else%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20span%20%3D%20torch.linspace(self.span%5B0%5D%2C%20self.span%5B1%5D%2C%20weight.shape%5B0%5D%20-%201)%0A%0A%20%20%20%20%20%20%20%20%20%20%20%20batch_size%20%3D%20x.shape%5B0%5D%0A%20%20%20%20%20%20%20%20%20%20%20%20assert%20(%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20x.shape%5B1%5D%20%2B%201%20%3D%3D%20weight.shape%5B0%5D%0A%20%20%20%20%20%20%20%20%20%20%20%20)%2C%20f%22%7Bx.shape%7D%20but%20weight%20shape%20%7Bweight.shape%7D%22%0A%0A%20%20%20%20%20%20%20%20%20%20%20%20%23%20Compute%20polynomial%20basis%3A%20(span%20-%20x)%5Ek%20for%20k%3D1..weight.shape%5B0%5D-1%0A%20%20%20%20%20%20%20%20%20%20%20%20powers%20%3D%20torch.arange(1%2C%20weight.shape%5B0%5D%2C%20dtype%3Dx.dtype%2C%20device%3Dx.device)%0A%20%20%20%20%20%20%20%20%20%20%20%20basis%20%3D%20(span%20-%20x)%20**%20powers%20%20%23%20%5Bbatch%2C%20num_powers%5D%0A%20%20%20%20%20%20%20%20%20%20%20%20x%20%3D%20weight%5B0%2C%20%3A%5D%20%2B%20basis%20%40%20weight%5B1%3A%2C%20%3A%5D%0A%20%20%20%20%20%20%20%20%20%20%20%20x%20%3D%20x.reshape(-1%2C%20weight.shape%5B1%5D)%0A%20%20%20%20%20%20%20%20%20%20%20%20assert%20(%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20x.shape%5B0%5D%20%3D%3D%20batch_size%0A%20%20%20%20%20%20%20%20%20%20%20%20)%2C%20f%22%7Bx.shape%7D%20but%20batch%20size%20%7Bbatch_size%7D%22%0A%20%20%20%20%20%20%20%20%20%20%20%20return%20x%0A%0A%20%20%20%20%20%20%20%20def%20_kann(self%2C%20x%2C%20weight)%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20k%20%3D%20weight.shape%5B0%5D%20-%201%0A%20%20%20%20%20%20%20%20%20%20%20%20span%20%3D%20torch.linspace(self.span%5B0%5D%2C%20self.span%5B1%5D%2C%20k%2C%20dtype%3Dx.dtype%2C%20device%3Dx.device)%0A%20%20%20%20%20%20%20%20%20%20%20%20batch_size%20%3D%20x.shape%5B0%5D%0A%20%20%20%20%20%20%20%20%20%20%20%20assert%20x.shape%5B1%5D%20%3D%3D%20weight.shape%5B1%5D%0A%0A%20%20%20%20%20%20%20%20%20%20%20%20%23%20Build%20polynomial%20basis%0A%20%20%20%20%20%20%20%20%20%20%20%20powers%20%3D%20torch.arange(1%2C%20k%20%2B%201%2C%20dtype%3Dx.dtype%2C%20device%3Dx.device)%0A%20%20%20%20%20%20%20%20%20%20%20%20ones%20%3D%20torch.ones(*x.shape%2C%201%2C%20dtype%3Dx.dtype%2C%20device%3Dx.device)%0A%20%20%20%20%20%20%20%20%20%20%20%20poly_basis%20%3D%20(span%20-%20x%5B%3A%2C%20%3A%2C%20None%5D)%20**%20powers%0A%20%20%20%20%20%20%20%20%20%20%20%20x%20%3D%20torch.cat(%5Bones%2C%20poly_basis%5D%2C%20dim%3D2)%0A%0A%20%20%20%20%20%20%20%20%20%20%20%20assert%20x.shape%5B1%5D%20%3D%3D%20weight.shape%5B1%5D%0A%20%20%20%20%20%20%20%20%20%20%20%20x%20%3D%20torch.einsum(%22bmk%2Ckmn-%3Ebn%22%2C%20x%2C%20weight)%0A%20%20%20%20%20%20%20%20%20%20%20%20x%20%3D%20x.reshape(-1%2C%20weight.shape%5B-1%5D)%0A%20%20%20%20%20%20%20%20%20%20%20%20assert%20(%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20x.shape%5B0%5D%20%3D%3D%20batch_size%0A%20%20%20%20%20%20%20%20%20%20%20%20)%2C%20f%22%7Bx.shape%7D%20but%20batch%20size%20%7Bbatch_size%7D%22%0A%20%20%20%20%20%20%20%20%20%20%20%20return%20x%0A%0A%20%20%20%20%20%20%20%20def%20forward(self%2C%20x)%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20for%20weight%20in%20self.weights%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20if%20weight.ndim%20%3D%3D%202%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20x%20%3D%20self._classic_equivalent(x%2C%20weight)%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20else%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20x%20%3D%20self._kann(x%2C%20weight)%0A%20%20%20%20%20%20%20%20%20%20%20%20return%20x%0A%0A%0A%20%20%20%20def%20build_classic_kan_model(%0A%20%20%20%20%20%20%20%20layers%3A%20list%5Bint%5D%20%3D%20None%2C%0A%20%20%20%20%20%20%20%20seed%3A%20int%20%3D%200%2C%0A%20%20%20%20%20%20%20%20output_dim%3A%20int%20%3D%201%2C%0A%20%20%20%20%20%20%20%20k%3A%20int%20%3D%201%2C%0A%20%20%20%20%20%20%20%20input_dim%3A%20int%20%3D%201%2C%0A%20%20%20%20%20%20%20%20span%3A%20tuple%5Bfloat%2C%20float%5D%20%3D%20(-1.0%2C%201.0)%2C%0A%20%20%20%20)%3A%0A%20%20%20%20%20%20%20%20if%20layers%20is%20None%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20layers%20%3D%20%5B%5D%0A%0A%20%20%20%20%20%20%20%20model%20%3D%20KAN(%0A%20%20%20%20%20%20%20%20%20%20%20%20layers%3Dlayers%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20input_dim%3Dinput_dim%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20output_dim%3Doutput_dim%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20k%3D-1%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20seed%3Dseed%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20span%3Dspan%2C%0A%20%20%20%20%20%20%20%20)%0A%20%20%20%20%20%20%20%20return%20model%0A%0A%0A%20%20%20%20def%20train_model(build_fn%2C%20x%2C%20y%2C%20layers%3DNone%2C%20num_epochs%3D1000%2C%20lr%3D0.01)%3A%0A%20%20%20%20%20%20%20%20model%20%3D%20build_fn(layers%3Dlayers)%0A%20%20%20%20%20%20%20%20optimizer%20%3D%20optim.Adam(model.parameters()%2C%20lr%3Dlr)%0A%20%20%20%20%20%20%20%20criterion%20%3D%20nn.MSELoss()%0A%0A%20%20%20%20%20%20%20%20%23%20Reshape%20y%20if%20needed%0A%20%20%20%20%20%20%20%20if%20y.ndim%20%3D%3D%201%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20y%20%3D%20y.unsqueeze(1)%0A%0A%20%20%20%20%20%20%20%20for%20epoch%20in%20range(num_epochs)%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%23%20Mini-batch%20training%0A%20%20%20%20%20%20%20%20%20%20%20%20for%20i%20in%20range(0%2C%20len(x)%2C%2020)%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20batch_x%20%3D%20x%5Bi%3Ai%2B20%5D%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20batch_y%20%3D%20y%5Bi%3Ai%2B20%5D%0A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20optimizer.zero_grad()%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20pred%20%3D%20model(batch_x)%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20loss%20%3D%20criterion(pred%2C%20batch_y)%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20loss.backward()%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20optimizer.step()%0A%0A%20%20%20%20%20%20%20%20return%20model%2C%20lambda%20inp%3A%20model(inp)%0A%0A%20%20%20%20return%20build_classic_kan_model%2C%20build_classic_nn_model%2C%20train_model%0A%0A%0A%40app.function%0A%23%20Generate%20some%20data%20and%20run%0Adef%20generate_data(fn%2C%20n%3D1000%2C%20seed%3D42%2C%20noise%3D0.01)%3A%0A%20%20%20%20torch.manual_seed(seed)%0A%20%20%20%20x%20%3D%20torch.linspace(-1%2C%201%2C%20n)%0A%20%20%20%20perm%20%3D%20torch.randperm(n)%0A%20%20%20%20x%20%3D%20x%5Bperm%5D%0A%20%20%20%20y%20%3D%20fn(x)%20%2B%20torch.randn(n)%20*%20noise%0A%20%20%20%20x%20%3D%20x.reshape(-1%2C%201)%0A%20%20%20%20return%20x%2C%20y%0A%0A%0A%40app.cell%0Adef%20_(build_classic_kan_model%2C%20build_classic_nn_model%2C%20train_model)%3A%0A%20%20%20%20%40mo.persistent_cache(method%3D%22lazy%22)%0A%20%20%20%20def%20kann_permutation(fn_idx%2C%20layers_idx%2C%20noise_idx)%3A%0A%20%20%20%20%20%20%20%20%23%20The%20single%20expensive%20computation%2C%20keyed%20purely%20on%20the%20discrete%20UI%0A%20%20%20%20%20%20%20%20%23%20indices.%20Everything%20torch-touching%20happens%20*inside*%20here%3B%20the%20return%0A%20%20%20%20%20%20%20%20%23%20value%20is%20torch-free%20(numpy%20%2B%20str)%2C%20so%20a%20cache%20hit%20restores%20it%20in%20the%0A%20%20%20%20%20%20%20%20%23%20browser%20without%20importing%20torch.%20The%20precompute%20sweep%20populates%20one%0A%20%20%20%20%20%20%20%20%23%20entry%20per%20(fn%2C%20layers%2C%20noise)%20index%3B%20the%20display%20cells%20below%20just%0A%20%20%20%20%20%20%20%20%23%20look%20up%20the%20entry%20for%20the%20current%20UI%20selection.%0A%20%20%20%20%20%20%20%20fn_label%2C%20fn_callable%20%3D%20FN_OPTIONS%5Bfn_idx%5D%0A%20%20%20%20%20%20%20%20layers_value%20%3D%20LAYER_OPTIONS%5Blayers_idx%5D%0A%0A%20%20%20%20%20%20%20%20x%2C%20y%20%3D%20generate_data(fn_callable%2C%20noise%3Dnoise_at(noise_idx))%0A%20%20%20%20%20%20%20%20classic_model%2C%20classic%20%3D%20train_model(%0A%20%20%20%20%20%20%20%20%20%20%20%20build_classic_nn_model%2C%20x%2C%20y%2C%20layers%3Dlayers_value%0A%20%20%20%20%20%20%20%20)%0A%20%20%20%20%20%20%20%20kan_model%2C%20kan%20%3D%20train_model(%0A%20%20%20%20%20%20%20%20%20%20%20%20build_classic_kan_model%2C%20x%2C%20y%2C%20layers%3Dlayers_value%0A%20%20%20%20%20%20%20%20)%0A%20%20%20%20%20%20%20%20return%20%7B%0A%20%20%20%20%20%20%20%20%20%20%20%20%22fn_label%22%3A%20fn_label%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%22layers_label%22%3A%20str(layers_value)%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%22x_np%22%3A%20x.numpy()%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%22y_np%22%3A%20y.numpy()%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%22classic_pred%22%3A%20classic(x).detach().numpy()%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%22kan_pred%22%3A%20kan(x).detach().numpy()%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%23%20Weights%20as%20numpy%20so%20the%20symbolic%20cell%20stays%20torch-free.%0A%20%20%20%20%20%20%20%20%20%20%20%20%22kan_weights%22%3A%20%5B%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20(name%2C%20weight.detach().numpy())%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20for%20name%2C%20weight%20in%20kan_model.named_parameters()%0A%20%20%20%20%20%20%20%20%20%20%20%20%5D%2C%0A%20%20%20%20%20%20%20%20%7D%0A%0A%20%20%20%20return%20(kann_permutation%2C)%0A%0A%0A%40app.cell%0Adef%20_(fn%2C%20kann_permutation%2C%20layers%2C%20noise)%3A%0A%20%20%20%20result%20%3D%20kann_permutation(fn.value%2C%20layers.value%2C%20noise.value)%0A%20%20%20%20return%20(result%2C)%0A%0A%0A%40app.cell%0Adef%20_(kann_permutation)%3A%0A%20%20%20%20from%20itertools%20import%20product%0A%0A%20%20%20%20%23%20Precompute%20sweep%3A%20warm%20the%20persistent%20cache%20for%20every%20point%20in%20the%0A%20%20%20%20%23%20discrete%20UI%20state%20space%20so%20the%20exported%20WASM%20notebook%20can%20serve%20any%0A%20%20%20%20%23%20selection%20torch-free.%20Skipped%20in%20the%20browser%20(torch%20is%20absent%20there%20and%0A%20%20%20%20%23%20the%20entries%20are%20already%20bundled)%3B%20runs%20natively%2C%20including%20inside%20the%0A%20%20%20%20%23%20%60export%20html-wasm%20--execute%60%20session%2C%20where%20the%20touched%20keys%20are%20what%0A%20%20%20%20%23%20gets%20bundled%20into%20%60public%2Fcache%2F%60.%20Run%20the%20precompute%20export%20against%20a%0A%20%20%20%20%23%20clean%20%60__marimo__%2Fcache%60%20so%20this%20cell%20isn't%20cache-skipped.%0A%20%20%20%20if%20not%20is_wasm()%3A%0A%20%20%20%20%20%20%20%20for%20_fn_idx%2C%20_layers_idx%2C%20_noise_idx%20in%20product(%0A%20%20%20%20%20%20%20%20%20%20%20%20range(len(FN_LABELS))%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20range(len(LAYER_OPTIONS))%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20range(NOISE_STEPS)%2C%0A%20%20%20%20%20%20%20%20)%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20kann_permutation(_fn_idx%2C%20_layers_idx%2C%20_noise_idx)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_()%3A%0A%20%20%20%20import%20sys%0A%20%20%20%20sys.version_info%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_()%3A%0A%20%20%20%20return%0A%0A%0Aif%20__name__%20%3D%3D%20%22__main__%22%3A%0A%20%20%20%20app.run()%0A