perf: use the permuted formulation

Author: avik-pal

bench/lux.jl

@@ -11,9 +11,9 @@ train!(args...; kwargs...) = train!(MSELoss(), AutoZygote(), args...; kwargs...)
 function train!(loss, backend, model, ps, st, data; epochs=10)
     l1 = loss(model, ps, st, first(data))
 
-    tstate = Lux.Experimental.TrainState(model, ps, st, Adam(0.01f0))
+    tstate = Training.TrainState(model, ps, st, Adam(0.01f0))
     for _ in 1:epochs, (x, y) in data
-        _, _, _, tstate = Lux.Experimental.single_train_step!(backend, loss, (x, y), tstate)
+        _, _, _, tstate = Training.single_train_step!(backend, loss, (x, y), tstate)
     end
 
     l2 = loss(model, ps, st, first(data))
@@ -25,14 +25,14 @@ end
 n_points = 128
 batch_size = 64
 
-x = rand(Float32, 1, n_points, batch_size);
-y = rand(Float32, 1, n_points, batch_size);
+x = rand(Float32, n_points, 1, batch_size);
+y = rand(Float32, n_points, 1, batch_size);
 data = [(x, y)];
 t_fwd = zeros(5)
 t_train = zeros(5)
 for i in 1:5
     chs = (1, 128, fill(64, i)..., 128, 1)
-    model = FourierNeuralOperator(gelu; chs=chs, modes=(16,))
+    model = FourierNeuralOperator(gelu; chs, modes=(16,), permuted=Val(true))
     ps, st = Lux.setup(rng, model)
     model(x, ps, st) # TTFX