feat: compile neural operators using Reactant

[avik-pal]

needs EnzymeAD/Reactant.jl#217 + a JLL bump

[avik-pal]

FNOs lead to crash of julia, need to investigate

[avik-pal]

some more progress. With EnzymeAD/Reactant.jl#245 the forward pass of FNO compiles. Though there seems to be an issue with the gradient EnzymeAD/Reactant.jl#246

[avik-pal]

Bypassed the rfft issue. But now getting EnzymeAD/Reactant.jl#238 (comment)

julia> ∇fno_compiled = @compile ∇fno(fno, ps, st, x)
error: 'complex.add' op operand #0 must be complex type with floating-point elements, but got 'tensor<16x64x64xcomplex<f32>>'
ERROR: "failed to run pass manager on module"
Stacktrace:
  [1] run!
    @ /mnt/software/lux/Reactant.jl/src/mlir/IR/Pass.jl:70 [inlined]
  [2] run_pass_pipeline!(mod::Reactant.MLIR.IR.Module, pass_pipeline::String)
    @ Reactant.Compiler /mnt/software/lux/Reactant.jl/src/Compiler.jl:251
  [3] compile_mlir!(mod::Reactant.MLIR.IR.Module, f::Function, args::Tuple{…}; optimize::Bool)
    @ Reactant.Compiler /mnt/software/lux/Reactant.jl/src/Compiler.jl:283
  [4] compile_mlir!
    @ /mnt/software/lux/Reactant.jl/src/Compiler.jl:265 [inlined]
  [5] (::Reactant.Compiler.var"#30#32"{Bool, typeof(∇fno), Tuple{…}})()
    @ Reactant.Compiler /mnt/software/lux/Reactant.jl/src/Compiler.jl:720
  [6] context!(f::Reactant.Compiler.var"#30#32"{Bool, typeof(∇fno), Tuple{…}}, ctx::Reactant.MLIR.IR.Context)
    @ Reactant.MLIR.IR /mnt/software/lux/Reactant.jl/src/mlir/IR/Context.jl:76
  [7] compile_xla(f::Function, args::Tuple{…}; client::Nothing, optimize::Bool)
    @ Reactant.Compiler /mnt/software/lux/Reactant.jl/src/Compiler.jl:717
  [8] compile_xla
    @ /mnt/software/lux/Reactant.jl/src/Compiler.jl:712 [inlined]
  [9] compile(f::Function, args::Tuple{…}; client::Nothing, optimize::Bool, sync::Bool)
    @ Reactant.Compiler /mnt/software/lux/Reactant.jl/src/Compiler.jl:744
 [10] top-level scope
    @ /mnt/software/lux/Reactant.jl/src/Compiler.jl:490
Some type information was truncated. Use `show(err)` to see complete types.

[wsmoses]

That should be resolved by the arith raise pass and converted to a stablehlo.add