DeepONet error with high dim input

[KirillZubov]

Either I didn't understand how or the implementation DeepOnet fail with high dim input data.

u = ones(10, 10,10)
y = ones(1, 10, 10)
deeponet = LuxNeuralOperators.DeepONet(
    Chain(Dense(10 => 10, Lux.tanh_fast), Dense(10 => 10, Lux.tanh_fast), Dense(10 => 10)),
    Chain(Dense(1 => 10, Lux.tanh_fast), Dense(10 => 10, Lux.tanh_fast),
        Dense(10 => 10, Lux.tanh_fast)))
ps, st = Lux.setup(Random.default_rng(), deeponet)

deeponet((u, y), ps, st)

ERROR: ArgumentError: no valid permutation of dimensions
Stacktrace:
 [1] permutedims(B::Array{Float64, 4}, perm::Tuple{Int64, Int64, Int64})
   @ Base ./multidimensional.jl:1597
 [2] #25
   @ ~/.julia/packages/Lux/TBg3E/src/helpers/compact.jl:357 [inlined]
 [3] (::LuxNeuralOperators.var"#25#26")(self#238::@NamedTuple{…}, ::Tuple{…}, ps#239::@NamedTuple{…}, st#240::@NamedTuple{…})
   @ LuxNeuralOperators ./none:0
 [4] macro expansion
   @ ~/.julia/packages/Lux/TBg3E/src/helpers/compact.jl:491 [inlined]
 [5] (::CompactLuxLayer{…})(x::Tuple{…}, ps::@NamedTuple{…}, st::@NamedTuple{…})
   @ Lux ~/.julia/packages/Lux/TBg3E/src/helpers/compact.jl:485
 [6] top-level scope
   @ REPL[867]:1
Some type information was truncated. Use `show(err)` to see complete types.

Probably problem here https://github.com/LuxDL/LuxNeuralOperators.jl/blob/67d90073f8fd18fdbeeb4475ece5189b7906dea8/src/deeponet.jl#L119-L122
it should be just dot product.