`Optimisers.jl` fails to compile

[avik-pal]

The last piece needed for LuxDL/Lux.jl#969

using Reactant, Optimisers

ps = (; a = rand(10), b = rand(10)) |> Reactant.to_rarray
st_opt = Optimisers.setup(Descent(0.1), ps)
dps = (; a = rand(10), b = rand(10)) |> Reactant.to_rarray

function update_params(ps, st_opt, dps)
    return Optimisers.update(st_opt, ps, dps)
end

@compile update_params(ps, st_opt, dps)

ERROR: "Unhandled type Type"
Stacktrace:
  [1] traced_type(::Type{Type}, seen::Tuple{Val{Core.TypeName}, Val{Core.TypeName}}, mode::Val{Reactant.TracedSetPath})
    @ Reactant /mnt/software/lux/Reactant.jl/src/Tracing.jl:103
  [2] traced_type(::Type{Core.TypeName}, seen::Tuple{}, mode::Val{Reactant.TracedSetPath})
    @ Reactant /mnt/software/lux/Reactant.jl/src/Tracing.jl:131
  [3] make_tracer(seen::Reactant.OrderedIdDict{…}, prev::Core.TypeName, path::Any, mode::Reactant.TraceMode; toscalar::Bool, tobatch::Nothing)
    @ Reactant /mnt/software/lux/Reactant.jl/src/Tracing.jl:238
  [4] make_tracer(seen::Reactant.OrderedIdDict{…}, prev::Type{…}, path::Any, mode::Reactant.TraceMode; toscalar::Bool, tobatch::Nothing)
    @ Reactant /mnt/software/lux/Reactant.jl/src/Tracing.jl:250
  [5] (::Reactant.var"#21#30"{Bool, Tuple{…}, Bool, Reactant.OrderedIdDict{…}})(i::Int64)
    @ Reactant /mnt/software/lux/Reactant.jl/src/utils.jl:32
  [6] ntuple
    @ ./ntuple.jl:19 [inlined]
  [7] make_mlir_fn(f::Function, args::Tuple{…}, kwargs::Tuple{}, name::String, concretein::Bool; toscalar::Bool)
    @ Reactant /mnt/software/lux/Reactant.jl/src/utils.jl:31
  [8] make_mlir_fn(f::Function, args::Tuple{…}, kwargs::Tuple{}, name::String, concretein::Bool)
    @ Reactant /mnt/software/lux/Reactant.jl/src/utils.jl:24
  [9] make_mlir_fn(f::Type, args::Tuple{…}, kwargs::Tuple{}, name::String, concretein::Bool; toscalar::Bool)
    @ Reactant /mnt/software/lux/Reactant.jl/src/utils.jl:26
 [10] make_mlir_fn
    @ /mnt/software/lux/Reactant.jl/src/utils.jl:24 [inlined]
 [11] elem_apply(f::Type, args::Reactant.TracedRArray{Float64, 1})
    @ Reactant /mnt/software/lux/Reactant.jl/src/TracedRArray.jl:344
 [12] _copyto!(dest::Reactant.TracedRArray{Float64, 1}, bc::Base.Broadcast.Broadcasted{Nothing, Tuple{…}, Type{…}, Tuple{…}})
    @ Reactant /mnt/software/lux/Reactant.jl/src/TracedRArray.jl:747
 [13] copyto!
    @ /mnt/software/lux/Reactant.jl/src/TracedRArray.jl:669 [inlined]
 [14] copyto!
    @ ./broadcast.jl:956 [inlined]
 [15] copy(bc::Base.Broadcast.Broadcasted{Reactant.AbstractReactantArrayStyle{…}, Tuple{…}, Type{…}, Tuple{…}})
    @ Reactant /mnt/software/lux/Reactant.jl/src/TracedRArray.jl:660
 [16] materialize
    @ ./broadcast.jl:903 [inlined]
 [17] subtract!(x::Reactant.TracedRArray{…}, x̄::Base.Broadcast.Broadcasted{…})
    @ Optimisers ~/.julia/packages/Optimisers/yDIWk/src/interface.jl:103
 [18] _update!(ℓ::Optimisers.Leaf{…}, x::Reactant.TracedRArray{…}; grads::IdDict{…}, params::IdDict{…})
    @ Optimisers ~/.julia/packages/Optimisers/yDIWk/src/interface.jl:97
 [19] _update!
    @ ~/.julia/packages/Optimisers/yDIWk/src/interface.jl:92 [inlined]
 [20] #8
    @ ~/.julia/packages/Optimisers/yDIWk/src/interface.jl:85 [inlined]
 [21] map
    @ ./tuple.jl:319 [inlined]
 [22] map
    @ ./namedtuple.jl:266 [inlined]
 [23] mapvalue
    @ ~/.julia/packages/Optimisers/yDIWk/src/utils.jl:2 [inlined]
 [24] _update!(tree::@NamedTuple{…}, x::@NamedTuple{…}; grads::IdDict{…}, params::IdDict{…})
    @ Optimisers ~/.julia/packages/Optimisers/yDIWk/src/interface.jl:85
 [25] _update!
    @ ~/.julia/packages/Optimisers/yDIWk/src/interface.jl:81 [inlined]
 [26] update!(::@NamedTuple{…}, ::@NamedTuple{…}, ::@NamedTuple{…})
    @ Optimisers ~/.julia/packages/Optimisers/yDIWk/src/interface.jl:77
 [27] update(::@NamedTuple{…}, ::@NamedTuple{…}, ::@NamedTuple{…})
    @ Optimisers ~/.julia/packages/Optimisers/yDIWk/src/interface.jl:67
 [28] update_params
    @ /mnt/software/lux/Reactant.jl/envs/optimisers/simple.jl:8 [inlined]
 [29] (::Tuple{})(none::@NamedTuple{…}, none::@NamedTuple{…}, none::@NamedTuple{…})
    @ Base.Experimental ./<missing>:0
 [30] (::Reactant.var"#26#35"{typeof(update_params), Reactant.MLIR.IR.Block, Vector{…}, Tuple{…}})()
    @ Reactant /mnt/software/lux/Reactant.jl/src/utils.jl:100
 [31] block!(f::Reactant.var"#26#35"{…}, blk::Reactant.MLIR.IR.Block)
    @ Reactant.MLIR.IR /mnt/software/lux/Reactant.jl/src/mlir/IR/Block.jl:201
 [32] make_mlir_fn(f::Function, args::Tuple{…}, kwargs::Tuple{}, name::String, concretein::Bool; toscalar::Bool)
    @ Reactant /mnt/software/lux/Reactant.jl/src/utils.jl:74
 [33] make_mlir_fn
    @ /mnt/software/lux/Reactant.jl/src/utils.jl:24 [inlined]
 [34] #6
    @ /mnt/software/lux/Reactant.jl/src/Compiler.jl:253 [inlined]
 [35] block!(f::Reactant.Compiler.var"#6#11"{typeof(update_params), Tuple{…}}, blk::Reactant.MLIR.IR.Block)
    @ Reactant.MLIR.IR /mnt/software/lux/Reactant.jl/src/mlir/IR/Block.jl:201
 [36] #5
    @ /mnt/software/lux/Reactant.jl/src/Compiler.jl:252 [inlined]
 [37] mmodule!(f::Reactant.Compiler.var"#5#10"{…}, blk::Reactant.MLIR.IR.Module)
    @ Reactant.MLIR.IR /mnt/software/lux/Reactant.jl/src/mlir/IR/Module.jl:93
 [38] compile_mlir!(mod::Reactant.MLIR.IR.Module, f::Function, args::Tuple{…}; optimize::Bool)
    @ Reactant.Compiler /mnt/software/lux/Reactant.jl/src/Compiler.jl:249
 [39] compile_mlir!
    @ /mnt/software/lux/Reactant.jl/src/Compiler.jl:248 [inlined]
 [40] (::Reactant.Compiler.var"#30#32"{typeof(update_params), Tuple{@NamedTuple{…}, @NamedTuple{…}, @NamedTuple{…}}})()
    @ Reactant.Compiler /mnt/software/lux/Reactant.jl/src/Compiler.jl:576
 [41] context!(f::Reactant.Compiler.var"#30#32"{typeof(update_params), Tuple{…}}, ctx::Reactant.MLIR.IR.Context)
    @ Reactant.MLIR.IR /mnt/software/lux/Reactant.jl/src/mlir/IR/Context.jl:71
 [42] compile_xla(f::Function, args::Tuple{@NamedTuple{…}, @NamedTuple{…}, @NamedTuple{…}}; client::Nothing)
    @ Reactant.Compiler /mnt/software/lux/Reactant.jl/src/Compiler.jl:573
 [43] compile_xla
    @ /mnt/software/lux/Reactant.jl/src/Compiler.jl:567 [inlined]
 [44] compile(f::Function, args::Tuple{@NamedTuple{…}, @NamedTuple{…}, @NamedTuple{…}}; client::Nothing)
    @ Reactant.Compiler /mnt/software/lux/Reactant.jl/src/Compiler.jl:600
 [45] compile(f::Function, args::Tuple{@NamedTuple{…}, @NamedTuple{…}, @NamedTuple{…}})
    @ Reactant.Compiler /mnt/software/lux/Reactant.jl/src/Compiler.jl:599
 [46] macro expansion
    @ /mnt/software/lux/Reactant.jl/src/Compiler.jl:360 [inlined]
 [47] top-level scope
    @ /mnt/software/lux/Reactant.jl/envs/optimisers/simple.jl:11
Some type information was truncated. Use `show(err)` to see complete types.