Struggling to use Mooncake (or Enzyme) backend

[gm89uk]

Hello all, Miles

I'm on Julia v1.11.5 but struggling to get mooncake to work (or Enzyme, although I'm not too sure of the syntax there).
I followed this example: #468 which worked when I copy pasted it. However, I cannot apply it to my code successfully.

using Distributed, CSV, DataFrames, SymbolicRegression, MLJBase
procs = addprocs(32)

@everywhere using SymbolicRegression, Mooncake, LoopVectorization, Bumper, DifferentiationInterface, Zygote

filepath = "....csv"
df       = CSV.read(filepath, DataFrame)
iol_map   = Dict(label => i for (i,label) in enumerate(unique(df.cts_Model)))
weightsIn = Float32.([iol_map[label] for label in df.cts_Model])
const n_val = maximum(Int.(weightsIn))
X = (
    sr_raw    = Float32.(df.sr_raw),
    cts_Model = Int32.(iol_map[l] for l in df.cts_Model)

)
y = Float32.(df.s)
@everywhere const n = $n_val
@everywhere begin
    expression_spec = @template_spec(expressions=(f,), parameters=(p1=n, p2=n)) do sr_raw, cts_Model
        return p1[cts_Model] * f(sr_raw) - p2[cts_Model]
    end
end
model = SRRegressor(
    elementwise_loss = L1DistLoss(),
    niterations          = 1_000_000,
    binary_operators     = [+, -, *, /],
    maxsize              = 5,
    bumper               = true,
    turbo                = true,
    parallelism          = :multiprocessing,
    procs                = procs,
    populations          = 32,
    population_size      = 100,
    parsimony            = 0.01,
    expression_spec      = expression_spec,
    autodiff_backend=AutoMooncake(; config=nothing), #if I comment this ln, the code runs fine
)
mach = machine(model, X, y)
fit!(mach)  

Error:

Some type information was truncated. Use `show(err)` to see complete types.

julia> Unhandled Task ERROR: EOFError: read end of file
Stacktrace:
 [1] read(this::Sockets.TCPSocket, ::Type{UInt8})
   @ Base .\stream.jl:1012
 [2] message_handler_loop(r_stream::Sockets.TCPSocket, w_stream::Sockets.TCPSocket, incoming::Bool)
   @ Distributed C:\Users\George\AppData\Local\Programs\Julia-1.11.5\share\julia\stdlib\v1.11\Distributed\src\process_messages.jl:183
 [3] process_tcp_streams(r_stream::Sockets.TCPSocket, w_stream::Sockets.TCPSocket, incoming::Bool)
   @ Distributed C:\Users\George\AppData\Local\Programs\Julia-1.11.5\share\julia\stdlib\v1.11\Distributed\src\process_messages.jl:133
 [4] (::Distributed.var"#103#104"{Sockets.TCPSocket, Sockets.TCPSocket, Bool})()
   @ Distributed C:\Users\George\AppData\Local\Programs\Julia-1.11.5\share\julia\stdlib\v1.11\Distributed\src\process_messages.jl:121

caused by: Error deserializing a remote exception from worker 4
Remote(original) exception of type CompositeException
Remote stacktrace :
Stacktrace:
  [1] threading_run
    @ .\threadingconstructs.jl:173
  [2] macro expansion
    @ .\threadingconstructs.jl:190 [inlined]
  [3] macro expansion
    @ C:\Users\George\.julia\packages\SymbolicRegression\fLqXe\src\Utils.jl:163 [inlined]
  [4] optimize_and_simplify_population
    @ C:\Users\George\.julia\packages\SymbolicRegression\fLqXe\src\SingleIteration.jl:79
  [5] macro expansion
    @ C:\Users\George\.julia\packages\SymbolicRegression\fLqXe\src\Configure.jl:354 [inlined]
  [6] #20
    @ C:\Users\George\AppData\Local\Programs\Julia-1.11.5\share\julia\stdlib\v1.11\Distributed\src\macros.jl:83
  [7] #invokelatest#2
    @ .\essentials.jl:1055 [inlined]
  [8] invokelatest
    @ .\essentials.jl:1052
  [9] #107
    @ C:\Users\George\AppData\Local\Programs\Julia-1.11.5\share\julia\stdlib\v1.11\Distributed\src\process_messages.jl:283
 [10] run_work_thunk
    @ C:\Users\George\AppData\Local\Programs\Julia-1.11.5\share\julia\stdlib\v1.11\Distributed\src\process_messages.jl:70
 [11] run_work_thunk
    @ C:\Users\George\AppData\Local\Programs\Julia-1.11.5\share\julia\stdlib\v1.11\Distributed\src\process_messages.jl:79
 [12] #100
    @ C:\Users\George\AppData\Local\Programs\Julia-1.11.5\share\julia\stdlib\v1.11\Distributed\src\process_messages.jl:88

...and 1 more exception.

Stacktrace:
  [1] deserialize(s::Distributed.ClusterSerializer{Sockets.TCPSocket}, t::Type{CapturedException})
    @ Distributed C:\Users\George\AppData\Local\Programs\Julia-1.11.5\share\julia\stdlib\v1.11\Distributed\src\clusterserialize.jl:226
  [2] handle_deserialize(s::Distributed.ClusterSerializer{Sockets.TCPSocket}, b::Int32)
    @ Serialization C:\Users\George\AppData\Local\Programs\Julia-1.11.5\share\julia\stdlib\v1.11\Serialization\src\Serialization.jl:908
  [3] deserialize(s::Distributed.ClusterSerializer{Sockets.TCPSocket}, t::DataType)
    @ Serialization C:\Users\George\AppData\Local\Programs\Julia-1.11.5\share\julia\stdlib\v1.11\Serialization\src\Serialization.jl:1585
  [4] handle_deserialize(s::Distributed.ClusterSerializer{Sockets.TCPSocket}, b::Int32)
    @ Serialization C:\Users\George\AppData\Local\Programs\Julia-1.11.5\share\julia\stdlib\v1.11\Serialization\src\Serialization.jl:908
  [5] deserialize
    @ C:\Users\George\AppData\Local\Programs\Julia-1.11.5\share\julia\stdlib\v1.11\Serialization\src\Serialization.jl:844 [inlined]
  [6] deserialize_msg(s::Distributed.ClusterSerializer{Sockets.TCPSocket})
    @ Distributed C:\Users\George\AppData\Local\Programs\Julia-1.11.5\share\julia\stdlib\v1.11\Distributed\src\messages.jl:87
  [7] #invokelatest#2
    @ .\essentials.jl:1055 [inlined]
  [8] invokelatest
    @ .\essentials.jl:1052 [inlined]
  [9] message_handler_loop(r_stream::Sockets.TCPSocket, w_stream::Sockets.TCPSocket, incoming::Bool)
    @ Distributed C:\Users\George\AppData\Local\Programs\Julia-1.11.5\share\julia\stdlib\v1.11\Distributed\src\process_messages.jl:176
 [10] process_tcp_streams(r_stream::Sockets.TCPSocket, w_stream::Sockets.TCPSocket, incoming::Bool)
    @ Distributed C:\Users\George\AppData\Local\Programs\Julia-1.11.5\share\julia\stdlib\v1.11\Distributed\src\process_messages.jl:133
 [11] (::Distributed.var"#103#104"{Sockets.TCPSocket, Sockets.TCPSocket, Bool})()
    @ Distributed C:\Users\George\AppData\Local\Programs\Julia-1.11.5\share\julia\stdlib\v1.11\Distributed\src\process_messages.jl:121

caused by: MethodError: no method matching memoryref(::MemoryRef{Float32}, ::Int64, ::Bool)
The function `memoryref` exists, but no method is defined for this combination of argument types.

Closest candidates are:
  memoryref(::GenericMemoryRef, ::Integer)
   @ Core boot.jl:524
  memoryref(::GenericMemory, ::Integer)
   @ Core boot.jl:523
  memoryref(::GenericMemory)
   @ Core boot.jl:522

Is there anything obvious I am missing? Thank you in advance.

[MilesCranmer]

I think Mooncake cannot (yet) differentiate through LoopVectorzation.jl (turbo) or Bumper.jl (bumper).

In any case this error could definitely be improved! I'm quite surprised by the stack trace

[gm89uk]

Thanks @MilesCranmer disabling them did the trick.
Could you also help me with the syntax for Enzyme and bumper/turbo compatibility?
I have 30 parameters in my categorical variable, so I was hoping using Enzyme / Mooncake could speed up template expressions, but I think bumper/turbo with default still edge out.

[MilesCranmer]

Yeah 30 parameters is right at the boundary of where reverse-mode autodiff starts to help. The default (ForwardDiff.jl) is likely just as fast! Especially since its compatible with turbo.

You can try Zygote too (also reverse mode)