Task switch error on Enzyme v0.13

[MilesCranmer]

When trying to use Enzyme as the autodiff backend for SymbolicRegression searches I ran into this error:

        nested task error: task switch not allowed from inside staged nor pure functions

The full stack trace:

┌ Error: Problem fitting the machine machine(SRRegressor(defaults = nothing, …), …). 
└ @ MLJBase ~/.julia/packages/MLJBase/7nGJF/src/machines.jl:694
[ Info: Running type checks... 
[ Info: Type checks okay. 
ERROR: LoadError: TaskFailedException
Stacktrace:
  [1] wait(t::Task)
    @ Base ./task.jl:370
  [2] fetch
    @ ./task.jl:390 [inlined]
  [3] _main_search_loop!(state::SymbolicRegression.SearchUtilsModule.SearchState{…}, datasets::Vector{…}, ropt::SymbolicRegression.SearchUtilsModule.RuntimeOptions{…}, options::Options{…})
    @ SymbolicRegression ~/PermaDocuments/SymbolicRegressionMonorepo/SymbolicRegression.jl/src/SymbolicRegression.jl:833
  [4] _equation_search(datasets::Vector{…}, ropt::SymbolicRegression.SearchUtilsModule.RuntimeOptions{…}, options::Options{…}, saved_state::Nothing)
    @ SymbolicRegression ~/PermaDocuments/SymbolicRegressionMonorepo/SymbolicRegression.jl/src/SymbolicRegression.jl:535
  [5] equation_search(datasets::Vector{…}; options::Options{…}, saved_state::Nothing, runtime_options::Nothing, runtime_options_kws::@Kwargs{…})
    @ SymbolicRegression ~/PermaDocuments/SymbolicRegressionMonorepo/SymbolicRegression.jl/src/SymbolicRegression.jl:525
  [6] equation_search(X::Matrix{…}, y::Matrix{…}; niterations::Int64, weights::Nothing, options::Options{…}, variable_names::Vector{…}, display_variable_names::Vector{…}, y_variable_names::Nothing, parallelism::Symbol, numprocs::Nothing, procs::Nothing, addprocs_function::Nothing, heap_size_hint_in_bytes::Nothing, runtests::Bool, saved_state::Nothing, return_state::Bool, run_id::Nothing, loss_type::Type{…}, verbosity::Int64, logger::Nothing, progress::Nothing, X_units::Nothing, y_units::Nothing, extra::@NamedTuple{…}, v_dim_out::Val{…}, multithreaded::Nothing)
    @ SymbolicRegression ~/PermaDocuments/SymbolicRegressionMonorepo/SymbolicRegression.jl/src/SymbolicRegression.jl:476
  [7] #equation_search#34
    @ ~/PermaDocuments/SymbolicRegressionMonorepo/SymbolicRegression.jl/src/SymbolicRegression.jl:499 [inlined]
  [8] _update(m::SRRegressor{…}, verbosity::Int64, old_fitresult::Nothing, old_cache::Nothing, X::@NamedTuple{…}, y::Vector{…}, w::Nothing, options::Options{…}, class::Vector{…})
    @ SymbolicRegression.MLJInterfaceModule ~/PermaDocuments/SymbolicRegressionMonorepo/SymbolicRegression.jl/src/MLJInterface.jl:253
  [9] _update(m::SRRegressor{…}, verbosity::Int64, old_fitresult::Nothing, old_cache::Nothing, X::@NamedTuple{…}, y::Vector{…}, w::Nothing, options::Options{…}, class::Nothing)
    @ SymbolicRegression.MLJInterfaceModule ~/PermaDocuments/SymbolicRegressionMonorepo/SymbolicRegression.jl/src/MLJInterface.jl:220
 [10] update(m::SRRegressor{DynamicQuantities.SymbolicDimensions{…}, DataType}, verbosity::Int64, old_fitresult::Nothing, old_cache::Nothing, X::@NamedTuple{x1::Vector{…}, x2::Vector{…}, class::Vector{…}}, y::Vector{Float64}, w::Nothing)
    @ SymbolicRegression.MLJInterfaceModule ~/PermaDocuments/SymbolicRegressionMonorepo/SymbolicRegression.jl/src/MLJInterface.jl:201
 [11] fit
    @ ~/PermaDocuments/SymbolicRegressionMonorepo/SymbolicRegression.jl/src/MLJInterface.jl:189 [inlined]
 [12] fit(m::SRRegressor{DynamicQuantities.SymbolicDimensions{DynamicQuantities.FixedRational{Int32, 25200}}, DataType}, verbosity::Int64, X::@NamedTuple{x1::Vector{Float64}, x2::Vector{Float64}, class::Vector{Int64}}, y::Vector{Float64})
    @ SymbolicRegression.MLJInterfaceModule ~/PermaDocuments/SymbolicRegressionMonorepo/SymbolicRegression.jl/src/MLJInterface.jl:189
 [13] fit_only!(mach::MLJBase.Machine{SRRegressor{DynamicQuantities.SymbolicDimensions{…}, DataType}, SRRegressor{DynamicQuantities.SymbolicDimensions{…}, DataType}, true}; rows::Nothing, verbosity::Int64, force::Bool, composite::Nothing)
    @ MLJBase ~/.julia/packages/MLJBase/7nGJF/src/machines.jl:692
 [14] fit_only!
    @ ~/.julia/packages/MLJBase/7nGJF/src/machines.jl:617 [inlined]
 [15] #fit!#63
    @ ~/.julia/packages/MLJBase/7nGJF/src/machines.jl:789 [inlined]
 [16] fit!(mach::MLJBase.Machine{SRRegressor{DynamicQuantities.SymbolicDimensions{DynamicQuantities.FixedRational{Int32, 25200}}, DataType}, SRRegressor{DynamicQuantities.SymbolicDimensions{DynamicQuantities.FixedRational{Int32, 25200}}, DataType}, true})
    @ MLJBase ~/.julia/packages/MLJBase/7nGJF/src/machines.jl:786
 [17] top-level scope
    @ ~/PermaDocuments/SymbolicRegressionMonorepo/SymbolicRegression.jl/examples/parameterized_function.jl:103
 [18] include(fname::String)
    @ Main ./sysimg.jl:38
 [19] top-level scope
    @ REPL[3]:1

    nested task error: TaskFailedException
    Stacktrace:
     [1] wait(t::Task)
       @ Base ./task.jl:370
     [2] fetch
       @ ./task.jl:390 [inlined]
     [3] (::SymbolicRegression.var"#89#94"{SymbolicRegression.SearchUtilsModule.SearchState{…}, Int64, Int64})()
       @ SymbolicRegression ~/PermaDocuments/SymbolicRegressionMonorepo/SymbolicRegression.jl/src/SymbolicRegression.jl:810
    
        nested task error: task switch not allowed from inside staged nor pure functions
        Stacktrace:
          [1] try_yieldto(undo::typeof(Base.ensure_rescheduled))
            @ Base ./task.jl:948
          [2] wait()
            @ Base ./task.jl:1022
          [3] wait(c::Base.GenericCondition{Base.Threads.SpinLock}; first::Bool)
            @ Base ./condition.jl:130
          [4] wait
            @ ./condition.jl:125 [inlined]
          [5] (::Base.var"#slowlock#733")(rl::ReentrantLock)
            @ Base ./lock.jl:157
          [6] lock
            @ ./lock.jl:147 [inlined]
          [7] cached_compilation
            @ ~/.julia/packages/Enzyme/RvNgp/src/compiler.jl:8412 [inlined]
          [8] thunkbase(ctx::LLVM.Context, mi::Core.MethodInstance, ::Val{…}, ::Type{…}, ::Type{…}, tt::Type{…}, ::Val{…}, ::Val{…}, ::Val{…}, ::Val{…}, ::Val{…}, ::Type{…}, ::Val{…}, ::Val{…})
            @ Enzyme.Compiler ~/.julia/packages/Enzyme/RvNgp/src/compiler.jl:8548
          [9] #s2104#19135
            @ ~/.julia/packages/Enzyme/RvNgp/src/compiler.jl:8685 [inlined]
         [10] 
            @ Enzyme.Compiler ./none:0
         [11] (::Core.GeneratedFunctionStub)(::UInt64, ::LineNumberNode, ::Any, ::Vararg{Any})
            @ Core ./boot.jl:707
         [12] autodiff
            @ ~/.julia/packages/Enzyme/RvNgp/src/Enzyme.jl:473 [inlined]
         [13] autodiff
            @ ~/.julia/packages/Enzyme/RvNgp/src/Enzyme.jl:537 [inlined]
         [14] autodiff
            @ ~/.julia/packages/Enzyme/RvNgp/src/Enzyme.jl:504 [inlined]
         [15] (::SymbolicRegression.ConstantOptimizationModule.GradEvaluator{…})(::Float64, G::Vector{…}, x::Vector{…})
            @ SymbolicRegressionEnzymeExt ~/PermaDocuments/SymbolicRegressionMonorepo/SymbolicRegression.jl/ext/SymbolicRegressionEnzymeExt.jl:42
         [16] (::NLSolversBase.var"#69#70"{NLSolversBase.InplaceObjective{…}, Float64})(G::Vector{Float64}, x::Vector{Float64})
            @ NLSolversBase ~/.julia/packages/NLSolversBase/kavn7/src/objective_types/incomplete.jl:54
         [17] value_gradient!!(obj::NLSolversBase.OnceDifferentiable{Float64, Vector{Float64}, Vector{Float64}}, x::Vector{Float64})
            @ NLSolversBase ~/.julia/packages/NLSolversBase/kavn7/src/interface.jl:82
         [18] initial_state(method::Optim.BFGS{…}, options::Optim.Options{…}, d::NLSolversBase.OnceDifferentiable{…}, initial_x::Vector{…})
            @ Optim ~/.julia/packages/Optim/ZhuZN/src/multivariate/solvers/first_order/bfgs.jl:94
         [19] optimize
            @ ~/.julia/packages/Optim/ZhuZN/src/multivariate/optimize/optimize.jl:36 [inlined]
         [20] optimize(f::NLSolversBase.InplaceObjective{…}, initial_x::Vector{…}, method::Optim.BFGS{…}, options::Optim.Options{…}; inplace::Bool, autodiff::Symbol)
            @ Optim ~/.julia/packages/Optim/ZhuZN/src/multivariate/optimize/interface.jl:143
         [21] optimize
            @ ~/.julia/packages/Optim/ZhuZN/src/multivariate/optimize/interface.jl:139 [inlined]
         [22] macro expansion
            @ ~/PermaDocuments/SymbolicRegressionMonorepo/SymbolicRegression.jl/src/ConstantOptimization.jl:76 [inlined]
         [23] _optimize_constants(dataset::Dataset{…}, member::PopMember{…}, options::Options{…}, algorithm::Optim.BFGS{…}, optimizer_options::Optim.Options{…}, idx::Nothing)
            @ SymbolicRegression.ConstantOptimizationModule ~/.julia/packages/DispatchDoctor/ZmxWH/src/stabilization.jl:314
         [24] macro expansion
            @ ~/PermaDocuments/SymbolicRegressionMonorepo/SymbolicRegression.jl/src/ConstantOptimization.jl:46 [inlined]
         [25] dispatch_optimize_constants(dataset::Dataset{…}, member::PopMember{…}, options::Options{…}, idx::Nothing)
            @ SymbolicRegression.ConstantOptimizationModule ~/.julia/packages/DispatchDoctor/ZmxWH/src/stabilization.jl:314
         [26] macro expansion
            @ ~/PermaDocuments/SymbolicRegressionMonorepo/SymbolicRegression.jl/src/ConstantOptimization.jl:27 [inlined]
         [27] optimize_constants(dataset::Dataset{…}, member::PopMember{…}, options::Options{…})
            @ SymbolicRegression.ConstantOptimizationModule ~/.julia/packages/DispatchDoctor/ZmxWH/src/stabilization.jl:314
         [28] macro expansion
            @ ~/PermaDocuments/SymbolicRegressionMonorepo/SymbolicRegression.jl/src/SingleIteration.jl:118 [inlined]
         [29] macro expansion
            @ ~/PermaDocuments/SymbolicRegressionMonorepo/SymbolicRegression.jl/src/Utils.jl:159 [inlined]
         [30] macro expansion
            @ ~/PermaDocuments/SymbolicRegressionMonorepo/SymbolicRegression.jl/src/SingleIteration.jl:109 [inlined]
         [31] optimize_and_simplify_population(dataset::Dataset{…}, pop::Population{…}, options::Options{…}, curmaxsize::Int64, record::Dict{…})
            @ SymbolicRegression.SingleIterationModule ~/.julia/packages/DispatchDoctor/ZmxWH/src/stabilization.jl:314
         [32] macro expansion
            @ ~/PermaDocuments/SymbolicRegressionMonorepo/SymbolicRegression.jl/src/SymbolicRegression.jl:1087 [inlined]
         [33] _dispatch_s_r_cycle(in_pop::Population{…}, dataset::Dataset{…}, options::Options{…}; pop::Int64, out::Int64, iteration::Int64, verbosity::Int64, cur_maxsize::Int64, running_search_statistics::SymbolicRegression.AdaptiveParsimonyModule.RunningSearchStatistics)
            @ SymbolicRegression ~/.julia/packages/DispatchDoctor/ZmxWH/src/stabilization.jl:314
         [34] macro expansion
            @ ~/PermaDocuments/SymbolicRegressionMonorepo/SymbolicRegression.jl/src/SymbolicRegression.jl:762 [inlined]
         [35] (::SymbolicRegression.var"#86#88"{…})()
            @ SymbolicRegression ~/PermaDocuments/SymbolicRegressionMonorepo/SymbolicRegression.jl/src/SearchUtils.jl:263
in expression starting at /Users/mcranmer/PermaDocuments/SymbolicRegressionMonorepo/SymbolicRegression.jl/examples/parameterized_function.jl:103
Some type information was truncated. Use `show(err)` to see complete types.

To reproduce, you can run the example here: https://ai.damtp.cam.ac.uk/symbolicregression/dev/examples/parameterized_function/ and swap :Zygote for :Enzyme.

For example:

using SymbolicRegression
using Random: MersenneTwister
using Enzyme
using MLJBase: machine, fit!, predict, report
using Test

X = let rng = MersenneTwister(0), n = 30
    (; x1=randn(rng, n), x2=randn(rng, n), class=rand(rng, 1:2, n))
end

y = let P1 = [0.1, 1.5], P2 = [3.2, 0.5]
    [2 * cos(x2 + P1[class]) + x1^2 - P2[class] for (x1, x2, class) in zip(X.x1, X.x2, X.class)]
end

model = SRRegressor(;
    niterations=100,
    binary_operators=[+, *, /, -],
    unary_operators=[cos, exp],
    populations=30,
    expression_type=ParametricExpression,
    expression_options=(; max_parameters=2),
    autodiff_backend=:Enzyme,
);

mach = machine(model, X, y)
fit!(mach)

Could it be because I am running Enzyme.jl from a task within the code?

More context: Enzyme.jl used to work for this, and I don't think I changed anything that would cause new behavior on my end. But I just switched to v0.13 so not sure if something changed. I can't test v0.12 due to the error here: #2080

[wsmoses]

@vchuravy this looks like a problem of julia not liking the locks of the caching mechanism all of a sudden??

I have no idea why

[wsmoses]

@MilesCranmer from the looks of it, it seems like this related to doing a first compile in a task

Separately, we've been consistently running the itnegration CI you made a while ago, and things continue to pass. Presumably that is set to the version at the time, so perhaps you can find what on your end changed to cause the issue?

[MilesCranmer]

The main difference is that the first compile on my system is in a worker thread, whereas in the CI, it’s the main thread. Nothing has changed on my side though.

[MilesCranmer]

Friendly ping on this

[wsmoses]

This issue is a bit out of my wheelhouse. @vchuravy's help will probably be needed here

[vchuravy]

@MilesCranmer it looks like you are calling autodiff from within the generator of a generated function?

[vchuravy]

Ah, no...

Enzyme uses a generated function (side-eyes billy) and from there invokes cached_compilation which takes a lock.
This might "happen" to work if it hits the fast-pass and there is no lock contention...

[vchuravy]

@MilesCranmer Can you try using Enzyme.set_abi(Reverse, NonGenABI) or Enzyme.set_abi(Forward, NonGenABI)?

[MilesCranmer]

(Just confirming – my call to Enzyme is indeed not within any generated functions)

• Enzyme.set_abi(Reverse, NonGenABI) => same error
• Enzyme.set_abi(Forward, NonGenABI) => same error
• [both] => same error

[wsmoses]

To be clear that’s not a global setting you would do

Enzyme.autodiff(set_abi(Forward, NonGenABI), …)

Can you also include the stack trace?

[MilesCranmer]

Oh I see. I only have reverse-mode set up at the moment; here's the full error with that tweak:

julia> julia --project=. examples/parameterized_function.jl
[ Info: Training machine(SRRegressor(defaults = nothing, …), …).
[ Info: Started!
┌ Error: Problem fitting the machine machine(SRRegressor(defaults = nothing, …), …). 
└ @ MLJBase ~/.julia/packages/MLJBase/7nGJF/src/machines.jl:694
[ Info: Running type checks... 
[ Info: Type checks okay. 
ERROR: LoadError: TaskFailedException
Stacktrace:
  [1] wait(t::Task)
    @ Base ./task.jl:370
  [2] fetch
    @ ./task.jl:390 [inlined]
  [3] _main_search_loop!(state::SymbolicRegression.SearchUtilsModule.SearchState{Float64, Float64, ParametricExpression{Float64, ParametricNode{Float64}, @NamedTuple{operators::Nothing, variable_names::Nothing, parameters::Matrix{Float64}, parameter_names::Nothing}}, Task, Channel}, datasets::Vector{Dataset{Float64, Float64, Matrix{Float64}, Vector{Float64}, Nothing, @NamedTuple{class::Vector{Int64}}, Nothing, Nothing, Nothing, Nothing}}, ropt::SymbolicRegression.SearchUtilsModule.RuntimeOptions{:multithreading, 1, true, Nothing}, options::Options{SymbolicRegression.CoreModule.OptionsStructModule.ComplexityMapping{Int64, Int64}, DynamicExpressions.OperatorEnumModule.OperatorEnum, ParametricNode, ParametricExpression, @NamedTuple{max_parameters::Int64}, MutationWeights, false, false, nothing, ADTypes.AutoEnzyme{Nothing, Nothing}, 5})
    @ SymbolicRegression ~/PermaDocuments/SymbolicRegressionMonorepo/SymbolicRegression.jl/src/SymbolicRegression.jl:833
  [4] _equation_search(datasets::Vector{Dataset{Float64, Float64, Matrix{Float64}, Vector{Float64}, Nothing, @NamedTuple{class::Vector{Int64}}, Nothing, Nothing, Nothing, Nothing}}, ropt::SymbolicRegression.SearchUtilsModule.RuntimeOptions{:multithreading, 1, true, Nothing}, options::Options{SymbolicRegression.CoreModule.OptionsStructModule.ComplexityMapping{Int64, Int64}, DynamicExpressions.OperatorEnumModule.OperatorEnum, ParametricNode, ParametricExpression, @NamedTuple{max_parameters::Int64}, MutationWeights, false, false, nothing, ADTypes.AutoEnzyme{Nothing, Nothing}, 5}, saved_state::Nothing)
    @ SymbolicRegression ~/PermaDocuments/SymbolicRegressionMonorepo/SymbolicRegression.jl/src/SymbolicRegression.jl:535
  [5] equation_search(datasets::Vector{Dataset{Float64, Float64, Matrix{Float64}, Vector{Float64}, Nothing, @NamedTuple{class::Vector{Int64}}, Nothing, Nothing, Nothing, Nothing}}; options::Options{SymbolicRegression.CoreModule.OptionsStructModule.ComplexityMapping{Int64, Int64}, DynamicExpressions.OperatorEnumModule.OperatorEnum, ParametricNode, ParametricExpression, @NamedTuple{max_parameters::Int64}, MutationWeights, false, false, nothing, ADTypes.AutoEnzyme{Nothing, Nothing}, 5}, saved_state::Nothing, runtime_options::Nothing, runtime_options_kws::@Kwargs{niterations::Int64, parallelism::Symbol, numprocs::Nothing, procs::Nothing, addprocs_function::Nothing, heap_size_hint_in_bytes::Nothing, runtests::Bool, return_state::Bool, run_id::Nothing, verbosity::Int64, logger::Nothing, progress::Nothing, v_dim_out::Val{1}})
    @ SymbolicRegression ~/PermaDocuments/SymbolicRegressionMonorepo/SymbolicRegression.jl/src/SymbolicRegression.jl:525
  [6] equation_search
    @ ~/PermaDocuments/SymbolicRegressionMonorepo/SymbolicRegression.jl/src/SymbolicRegression.jl:506 [inlined]
  [7] equation_search(X::Matrix{Float64}, y::Matrix{Float64}; niterations::Int64, weights::Nothing, options::Options{SymbolicRegression.CoreModule.OptionsStructModule.ComplexityMapping{Int64, Int64}, DynamicExpressions.OperatorEnumModule.OperatorEnum, ParametricNode, ParametricExpression, @NamedTuple{max_parameters::Int64}, MutationWeights, false, false, nothing, ADTypes.AutoEnzyme{Nothing, Nothing}, 5}, variable_names::Vector{String}, display_variable_names::Vector{String}, y_variable_names::Nothing, parallelism::Symbol, numprocs::Nothing, procs::Nothing, addprocs_function::Nothing, heap_size_hint_in_bytes::Nothing, runtests::Bool, saved_state::Nothing, return_state::Bool, run_id::Nothing, loss_type::Type{Nothing}, verbosity::Int64, logger::Nothing, progress::Nothing, X_units::Nothing, y_units::Nothing, extra::@NamedTuple{class::Vector{Int64}}, v_dim_out::Val{1}, multithreaded::Nothing)
    @ SymbolicRegression ~/PermaDocuments/SymbolicRegressionMonorepo/SymbolicRegression.jl/src/SymbolicRegression.jl:476
  [8] #equation_search#21
    @ ~/PermaDocuments/SymbolicRegressionMonorepo/SymbolicRegression.jl/src/SymbolicRegression.jl:499 [inlined]
  [9] _update(m::SRRegressor{DynamicQuantities.SymbolicDimensions{DynamicQuantities.FixedRational{Int32, 25200}}, DataType}, verbosity::Int64, old_fitresult::Nothing, old_cache::Nothing, X::@NamedTuple{x1::Vector{Float64}, x2::Vector{Float64}}, y::Vector{Float64}, w::Nothing, options::Options{SymbolicRegression.CoreModule.OptionsStructModule.ComplexityMapping{Int64, Int64}, DynamicExpressions.OperatorEnumModule.OperatorEnum, ParametricNode, ParametricExpression, @NamedTuple{max_parameters::Int64}, MutationWeights, false, false, nothing, ADTypes.AutoEnzyme{Nothing, Nothing}, 5}, class::Vector{Int64})
    @ SymbolicRegression.MLJInterfaceModule ~/PermaDocuments/SymbolicRegressionMonorepo/SymbolicRegression.jl/src/MLJInterface.jl:253
 [10] _update(m::SRRegressor{DynamicQuantities.SymbolicDimensions{DynamicQuantities.FixedRational{Int32, 25200}}, DataType}, verbosity::Int64, old_fitresult::Nothing, old_cache::Nothing, X::@NamedTuple{x1::Vector{Float64}, x2::Vector{Float64}, class::Vector{Int64}}, y::Vector{Float64}, w::Nothing, options::Options{SymbolicRegression.CoreModule.OptionsStructModule.ComplexityMapping{Int64, Int64}, DynamicExpressions.OperatorEnumModule.OperatorEnum, ParametricNode, ParametricExpression, @NamedTuple{max_parameters::Int64}, MutationWeights, false, false, nothing, ADTypes.AutoEnzyme{Nothing, Nothing}, 5}, class::Nothing)
    @ SymbolicRegression.MLJInterfaceModule ~/PermaDocuments/SymbolicRegressionMonorepo/SymbolicRegression.jl/src/MLJInterface.jl:220
 [11] update(m::SRRegressor{DynamicQuantities.SymbolicDimensions{DynamicQuantities.FixedRational{Int32, 25200}}, DataType}, verbosity::Int64, old_fitresult::Nothing, old_cache::Nothing, X::@NamedTuple{x1::Vector{Float64}, x2::Vector{Float64}, class::Vector{Int64}}, y::Vector{Float64}, w::Nothing)
    @ SymbolicRegression.MLJInterfaceModule ~/PermaDocuments/SymbolicRegressionMonorepo/SymbolicRegression.jl/src/MLJInterface.jl:201
 [12] fit
    @ ~/PermaDocuments/SymbolicRegressionMonorepo/SymbolicRegression.jl/src/MLJInterface.jl:189 [inlined]
 [13] fit(m::SRRegressor{DynamicQuantities.SymbolicDimensions{DynamicQuantities.FixedRational{Int32, 25200}}, DataType}, verbosity::Int64, X::@NamedTuple{x1::Vector{Float64}, x2::Vector{Float64}, class::Vector{Int64}}, y::Vector{Float64})
    @ SymbolicRegression.MLJInterfaceModule ~/PermaDocuments/SymbolicRegressionMonorepo/SymbolicRegression.jl/src/MLJInterface.jl:189
 [14] fit_only!(mach::MLJBase.Machine{SRRegressor{DynamicQuantities.SymbolicDimensions{DynamicQuantities.FixedRational{Int32, 25200}}, DataType}, SRRegressor{DynamicQuantities.SymbolicDimensions{DynamicQuantities.FixedRational{Int32, 25200}}, DataType}, true}; rows::Nothing, verbosity::Int64, force::Bool, composite::Nothing)
    @ MLJBase ~/.julia/packages/MLJBase/7nGJF/src/machines.jl:692
 [15] fit_only!
    @ ~/.julia/packages/MLJBase/7nGJF/src/machines.jl:617 [inlined]
 [16] #fit!#63
    @ ~/.julia/packages/MLJBase/7nGJF/src/machines.jl:789 [inlined]
 [17] fit!(mach::MLJBase.Machine{SRRegressor{DynamicQuantities.SymbolicDimensions{DynamicQuantities.FixedRational{Int32, 25200}}, DataType}, SRRegressor{DynamicQuantities.SymbolicDimensions{DynamicQuantities.FixedRational{Int32, 25200}}, DataType}, true})
    @ MLJBase ~/.julia/packages/MLJBase/7nGJF/src/machines.jl:786
 [18] top-level scope
    @ ~/PermaDocuments/SymbolicRegressionMonorepo/SymbolicRegression.jl/examples/parameterized_function.jl:101

    nested task error: TaskFailedException
    Stacktrace:
     [1] wait(t::Task)
       @ Base ./task.jl:370
     [2] fetch
       @ ./task.jl:390 [inlined]
     [3] (::SymbolicRegression.var"#56#61"{SymbolicRegression.SearchUtilsModule.SearchState{Float64, Float64, ParametricExpression{Float64, ParametricNode{Float64}, @NamedTuple{operators::Nothing, variable_names::Nothing, parameters::Matrix{Float64}, parameter_names::Nothing}}, Task, Channel}, Int64, Int64})()
       @ SymbolicRegression ~/PermaDocuments/SymbolicRegressionMonorepo/SymbolicRegression.jl/src/SymbolicRegression.jl:810
    
        nested task error: TaskFailedException
        Stacktrace:
          [1] wait(t::Task)
            @ Base ./task.jl:370
          [2] fetch
            @ ./task.jl:390 [inlined]
          [3] with_stacksize
            @ ~/PermaDocuments/SymbolicRegressionMonorepo/SymbolicRegression.jl/ext/SymbolicRegressionEnzymeExt.jl:31 [inlined]
          [4] (::SymbolicRegression.ConstantOptimizationModule.GradEvaluator{SymbolicRegression.ConstantOptimizationModule.Evaluator{ParametricExpression{Float64, ParametricNode{Float64}, @NamedTuple{operators::Nothing, variable_names::Nothing, parameters::Matrix{Float64}, parameter_names::Nothing}}, @NamedTuple{constant_refs::Vector{Base.RefValue{ParametricNode{Float64}}}, parameter_refs::Matrix{Float64}, num_parameters::Int64, num_constants::Int64}, Dataset{Float64, Float64, Matrix{Float64}, Vector{Float64}, Nothing, @NamedTuple{class::Vector{Int64}}, Nothing, Nothing, Nothing, Nothing}, Options{SymbolicRegression.CoreModule.OptionsStructModule.ComplexityMapping{Int64, Int64}, DynamicExpressions.OperatorEnumModule.OperatorEnum{Tuple{typeof(+), typeof(*), typeof(/), typeof(-)}, Tuple{typeof(cos), typeof(exp)}}, ParametricNode, ParametricExpression, @NamedTuple{max_parameters::Int64}, MutationWeights, false, false, nothing, ADTypes.AutoEnzyme{Nothing, Nothing}, 5}, Nothing}, ADTypes.AutoEnzyme{Nothing, Nothing}, @NamedTuple{storage_tree::ParametricExpression{Float64, ParametricNode{Float64}, @NamedTuple{operators::Nothing, variable_names::Nothing, parameters::Matrix{Float64}, parameter_names::Nothing}}, storage_refs::@NamedTuple{constant_refs::Vector{Base.RefValue{ParametricNode{Float64}}}, parameter_refs::Matrix{Float64}, num_parameters::Int64, num_constants::Int64}, storage_dataset::Dataset{Float64, Float64, Matrix{Float64}, Vector{Float64}, Nothing, @NamedTuple{class::Vector{Int64}}, Nothing, Nothing, Nothing, Nothing}}})(::Float64, G::Vector{Float64}, x::Vector{Float64})
            @ SymbolicRegressionEnzymeExt ~/PermaDocuments/SymbolicRegressionMonorepo/SymbolicRegression.jl/ext/SymbolicRegressionEnzymeExt.jl:41
          [5] (::NLSolversBase.var"#69#70"{NLSolversBase.InplaceObjective{Nothing, SymbolicRegression.ConstantOptimizationModule.GradEvaluator{SymbolicRegression.ConstantOptimizationModule.Evaluator{ParametricExpression{Float64, ParametricNode{Float64}, @NamedTuple{operators::Nothing, variable_names::Nothing, parameters::Matrix{Float64}, parameter_names::Nothing}}, @NamedTuple{constant_refs::Vector{Base.RefValue{ParametricNode{Float64}}}, parameter_refs::Matrix{Float64}, num_parameters::Int64, num_constants::Int64}, Dataset{Float64, Float64, Matrix{Float64}, Vector{Float64}, Nothing, @NamedTuple{class::Vector{Int64}}, Nothing, Nothing, Nothing, Nothing}, Options{SymbolicRegression.CoreModule.OptionsStructModule.ComplexityMapping{Int64, Int64}, DynamicExpressions.OperatorEnumModule.OperatorEnum{Tuple{typeof(+), typeof(*), typeof(/), typeof(-)}, Tuple{typeof(cos), typeof(exp)}}, ParametricNode, ParametricExpression, @NamedTuple{max_parameters::Int64}, MutationWeights, false, false, nothing, ADTypes.AutoEnzyme{Nothing, Nothing}, 5}, Nothing}, ADTypes.AutoEnzyme{Nothing, Nothing}, @NamedTuple{storage_tree::ParametricExpression{Float64, ParametricNode{Float64}, @NamedTuple{operators::Nothing, variable_names::Nothing, parameters::Matrix{Float64}, parameter_names::Nothing}}, storage_refs::@NamedTuple{constant_refs::Vector{Base.RefValue{ParametricNode{Float64}}}, parameter_refs::Matrix{Float64}, num_parameters::Int64, num_constants::Int64}, storage_dataset::Dataset{Float64, Float64, Matrix{Float64}, Vector{Float64}, Nothing, @NamedTuple{class::Vector{Int64}}, Nothing, Nothing, Nothing, Nothing}}}, Nothing, Nothing, Nothing}, Float64})(G::Vector{Float64}, x::Vector{Float64})
            @ NLSolversBase ~/.julia/packages/NLSolversBase/kavn7/src/objective_types/incomplete.jl:54
          [6] value_gradient!!(obj::NLSolversBase.OnceDifferentiable{Float64, Vector{Float64}, Vector{Float64}}, x::Vector{Float64})
            @ NLSolversBase ~/.julia/packages/NLSolversBase/kavn7/src/interface.jl:82
          [7] initial_state(method::Optim.BFGS{LineSearches.InitialStatic{Float64}, LineSearches.BackTracking{Float64, Int64}, Nothing, Nothing, Optim.Flat}, options::Optim.Options{Float64, Nothing}, d::NLSolversBase.OnceDifferentiable{Float64, Vector{Float64}, Vector{Float64}}, initial_x::Vector{Float64})
            @ Optim ~/.julia/packages/Optim/fBdaz/src/multivariate/solvers/first_order/bfgs.jl:94
          [8] optimize
            @ ~/.julia/packages/Optim/fBdaz/src/multivariate/optimize/optimize.jl:36 [inlined]
          [9] optimize(f::NLSolversBase.InplaceObjective{Nothing, SymbolicRegression.ConstantOptimizationModule.GradEvaluator{SymbolicRegression.ConstantOptimizationModule.Evaluator{ParametricExpression{Float64, ParametricNode{Float64}, @NamedTuple{operators::Nothing, variable_names::Nothing, parameters::Matrix{Float64}, parameter_names::Nothing}}, @NamedTuple{constant_refs::Vector{Base.RefValue{ParametricNode{Float64}}}, parameter_refs::Matrix{Float64}, num_parameters::Int64, num_constants::Int64}, Dataset{Float64, Float64, Matrix{Float64}, Vector{Float64}, Nothing, @NamedTuple{class::Vector{Int64}}, Nothing, Nothing, Nothing, Nothing}, Options{SymbolicRegression.CoreModule.OptionsStructModule.ComplexityMapping{Int64, Int64}, DynamicExpressions.OperatorEnumModule.OperatorEnum{Tuple{typeof(+), typeof(*), typeof(/), typeof(-)}, Tuple{typeof(cos), typeof(exp)}}, ParametricNode, ParametricExpression, @NamedTuple{max_parameters::Int64}, MutationWeights, false, false, nothing, ADTypes.AutoEnzyme{Nothing, Nothing}, 5}, Nothing}, ADTypes.AutoEnzyme{Nothing, Nothing}, @NamedTuple{storage_tree::ParametricExpression{Float64, ParametricNode{Float64}, @NamedTuple{operators::Nothing, variable_names::Nothing, parameters::Matrix{Float64}, parameter_names::Nothing}}, storage_refs::@NamedTuple{constant_refs::Vector{Base.RefValue{ParametricNode{Float64}}}, parameter_refs::Matrix{Float64}, num_parameters::Int64, num_constants::Int64}, storage_dataset::Dataset{Float64, Float64, Matrix{Float64}, Vector{Float64}, Nothing, @NamedTuple{class::Vector{Int64}}, Nothing, Nothing, Nothing, Nothing}}}, Nothing, Nothing, Nothing}, initial_x::Vector{Float64}, method::Optim.BFGS{LineSearches.InitialStatic{Float64}, LineSearches.BackTracking{Float64, Int64}, Nothing, Nothing, Optim.Flat}, options::Optim.Options{Float64, Nothing}; inplace::Bool, autodiff::Symbol)
            @ Optim ~/.julia/packages/Optim/fBdaz/src/multivariate/optimize/interface.jl:143
         [10] optimize
            @ ~/.julia/packages/Optim/fBdaz/src/multivariate/optimize/interface.jl:139 [inlined]
         [11] _optimize_constants(dataset::Dataset{Float64, Float64, Matrix{Float64}, Vector{Float64}, Nothing, @NamedTuple{class::Vector{Int64}}, Nothing, Nothing, Nothing, Nothing}, member::PopMember{Float64, Float64, ParametricExpression{Float64, ParametricNode{Float64}, @NamedTuple{operators::Nothing, variable_names::Nothing, parameters::Matrix{Float64}, parameter_names::Nothing}}}, options::Options{SymbolicRegression.CoreModule.OptionsStructModule.ComplexityMapping{Int64, Int64}, DynamicExpressions.OperatorEnumModule.OperatorEnum{Tuple{typeof(+), typeof(*), typeof(/), typeof(-)}, Tuple{typeof(cos), typeof(exp)}}, ParametricNode, ParametricExpression, @NamedTuple{max_parameters::Int64}, MutationWeights, false, false, nothing, ADTypes.AutoEnzyme{Nothing, Nothing}, 5}, algorithm::Optim.BFGS{LineSearches.InitialStatic{Float64}, LineSearches.BackTracking{Float64, Int64}, Nothing, Nothing, Optim.Flat}, optimizer_options::Optim.Options{Float64, Nothing}, idx::Nothing)
            @ SymbolicRegression.ConstantOptimizationModule ~/PermaDocuments/SymbolicRegressionMonorepo/SymbolicRegression.jl/src/ConstantOptimization.jl:76
         [12] dispatch_optimize_constants(dataset::Dataset{Float64, Float64, Matrix{Float64}, Vector{Float64}, Nothing, @NamedTuple{class::Vector{Int64}}, Nothing, Nothing, Nothing, Nothing}, member::PopMember{Float64, Float64, ParametricExpression{Float64, ParametricNode{Float64}, @NamedTuple{operators::Nothing, variable_names::Nothing, parameters::Matrix{Float64}, parameter_names::Nothing}}}, options::Options{SymbolicRegression.CoreModule.OptionsStructModule.ComplexityMapping{Int64, Int64}, DynamicExpressions.OperatorEnumModule.OperatorEnum, ParametricNode, ParametricExpression, @NamedTuple{max_parameters::Int64}, MutationWeights, false, false, nothing, ADTypes.AutoEnzyme{Nothing, Nothing}, 5}, idx::Nothing)
            @ SymbolicRegression.ConstantOptimizationModule ~/PermaDocuments/SymbolicRegressionMonorepo/SymbolicRegression.jl/src/ConstantOptimization.jl:46
         [13] optimize_constants(dataset::Dataset{Float64, Float64, Matrix{Float64}, Vector{Float64}, Nothing, @NamedTuple{class::Vector{Int64}}, Nothing, Nothing, Nothing, Nothing}, member::PopMember{Float64, Float64, ParametricExpression{Float64, ParametricNode{Float64}, @NamedTuple{operators::Nothing, variable_names::Nothing, parameters::Matrix{Float64}, parameter_names::Nothing}}}, options::Options{SymbolicRegression.CoreModule.OptionsStructModule.ComplexityMapping{Int64, Int64}, DynamicExpressions.OperatorEnumModule.OperatorEnum, ParametricNode, ParametricExpression, @NamedTuple{max_parameters::Int64}, MutationWeights, false, false, nothing, ADTypes.AutoEnzyme{Nothing, Nothing}, 5})
            @ SymbolicRegression.ConstantOptimizationModule ~/PermaDocuments/SymbolicRegressionMonorepo/SymbolicRegression.jl/src/ConstantOptimization.jl:27
         [14] macro expansion
            @ ~/PermaDocuments/SymbolicRegressionMonorepo/SymbolicRegression.jl/src/SingleIteration.jl:118 [inlined]
         [15] macro expansion
            @ ~/PermaDocuments/SymbolicRegressionMonorepo/SymbolicRegression.jl/src/Utils.jl:159 [inlined]
         [16] optimize_and_simplify_population(dataset::Dataset{Float64, Float64, Matrix{Float64}, Vector{Float64}, Nothing, @NamedTuple{class::Vector{Int64}}, Nothing, Nothing, Nothing, Nothing}, pop::Population{Float64, Float64, ParametricExpression{Float64, ParametricNode{Float64}, @NamedTuple{operators::Nothing, variable_names::Nothing, parameters::Matrix{Float64}, parameter_names::Nothing}}}, options::Options{SymbolicRegression.CoreModule.OptionsStructModule.ComplexityMapping{Int64, Int64}, DynamicExpressions.OperatorEnumModule.OperatorEnum, ParametricNode, ParametricExpression, @NamedTuple{max_parameters::Int64}, MutationWeights, false, false, nothing, ADTypes.AutoEnzyme{Nothing, Nothing}, 5}, curmaxsize::Int64, record::Dict{String, Any})
            @ SymbolicRegression.SingleIterationModule ~/PermaDocuments/SymbolicRegressionMonorepo/SymbolicRegression.jl/src/SingleIteration.jl:109
         [17] _dispatch_s_r_cycle(in_pop::Population{Float64, Float64, ParametricExpression{Float64, ParametricNode{Float64}, @NamedTuple{operators::Nothing, variable_names::Nothing, parameters::Matrix{Float64}, parameter_names::Nothing}}}, dataset::Dataset{Float64, Float64, Matrix{Float64}, Vector{Float64}, Nothing, @NamedTuple{class::Vector{Int64}}, Nothing, Nothing, Nothing, Nothing}, options::Options{SymbolicRegression.CoreModule.OptionsStructModule.ComplexityMapping{Int64, Int64}, DynamicExpressions.OperatorEnumModule.OperatorEnum, ParametricNode, ParametricExpression, @NamedTuple{max_parameters::Int64}, MutationWeights, false, false, nothing, ADTypes.AutoEnzyme{Nothing, Nothing}, 5}; pop::Int64, out::Int64, iteration::Int64, verbosity::Int64, cur_maxsize::Int64, running_search_statistics::SymbolicRegression.AdaptiveParsimonyModule.RunningSearchStatistics)
            @ SymbolicRegression ~/PermaDocuments/SymbolicRegressionMonorepo/SymbolicRegression.jl/src/SymbolicRegression.jl:1087
         [18] macro expansion
            @ ~/PermaDocuments/SymbolicRegressionMonorepo/SymbolicRegression.jl/src/SymbolicRegression.jl:762 [inlined]
         [19] (::SymbolicRegression.var"#53#55"{Float64, ParametricExpression{Float64, ParametricNode{Float64}, @NamedTuple{operators::Nothing, variable_names::Nothing, parameters::Matrix{Float64}, parameter_names::Nothing}}, Float64, SymbolicRegression.SearchUtilsModule.RuntimeOptions{:multithreading, 1, true, Nothing}, Options{SymbolicRegression.CoreModule.OptionsStructModule.ComplexityMapping{Int64, Int64}, DynamicExpressions.OperatorEnumModule.OperatorEnum, ParametricNode, ParametricExpression, @NamedTuple{max_parameters::Int64}, MutationWeights, false, false, nothing, ADTypes.AutoEnzyme{Nothing, Nothing}, 5}, Int64, Task, SymbolicRegression.AdaptiveParsimonyModule.RunningSearchStatistics, Int64, Dataset{Float64, Float64, Matrix{Float64}, Vector{Float64}, Nothing, @NamedTuple{class::Vector{Int64}}, Nothing, Nothing, Nothing, Nothing}, Int64})()
            @ SymbolicRegression ~/PermaDocuments/SymbolicRegressionMonorepo/SymbolicRegression.jl/src/SearchUtils.jl:263
        
            nested task error: AssertionError: Base.isconcretetype(typ)
            Stacktrace:
              [1] abs_typeof(arg::LLVM.ExtractValueInst, partial::Bool, seenphis::Set{LLVM.PHIInst})
                @ Enzyme.Compiler ~/.julia/packages/Enzyme/RvNgp/src/absint.jl:614
              [2] abs_typeof(arg::LLVM.ExtractValueInst, partial::Bool)
                @ Enzyme.Compiler ~/.julia/packages/Enzyme/RvNgp/src/absint.jl:281
              [3] codegen(output::Symbol, job::GPUCompiler.CompilerJob{Enzyme.Compiler.EnzymeTarget, Enzyme.Compiler.EnzymeCompilerParams}; libraries::Bool, deferred_codegen::Bool, optimize::Bool, toplevel::Bool, strip::Bool, validate::Bool, only_entry::Bool, parent_job::Nothing)
                @ Enzyme.Compiler ~/.julia/packages/Enzyme/RvNgp/src/compiler.jl:7095
              [4] codegen
                @ ~/.julia/packages/Enzyme/RvNgp/src/compiler.jl:6072 [inlined]
              [5] _thunk(job::GPUCompiler.CompilerJob{Enzyme.Compiler.EnzymeTarget, Enzyme.Compiler.EnzymeCompilerParams}, postopt::Bool)
                @ Enzyme.Compiler ~/.julia/packages/Enzyme/RvNgp/src/compiler.jl:8375
              [6] cached_compilation(job::GPUCompiler.CompilerJob)
                @ Enzyme.Compiler ~/.julia/packages/Enzyme/RvNgp/src/compiler.jl:8416
              [7] thunkbase
                @ ~/.julia/packages/Enzyme/RvNgp/src/compiler.jl:8548 [inlined]
              [8] thunk
                @ ~/.julia/packages/Enzyme/RvNgp/src/compiler.jl:8631 [inlined]
              [9] autodiff
                @ ~/.julia/packages/Enzyme/RvNgp/src/Enzyme.jl:473 [inlined]
             [10] autodiff
                @ ~/.julia/packages/Enzyme/RvNgp/src/Enzyme.jl:537 [inlined]
             [11] autodiff
                @ ~/.julia/packages/Enzyme/RvNgp/src/Enzyme.jl:504 [inlined]
             [12] (::SymbolicRegressionEnzymeExt.var"#1#2"{SymbolicRegression.ConstantOptimizationModule.GradEvaluator{SymbolicRegression.ConstantOptimizationModule.Evaluator{ParametricExpression{Float64, ParametricNode{Float64}, @NamedTuple{operators::Nothing, variable_names::Nothing, parameters::Matrix{Float64}, parameter_names::Nothing}}, @NamedTuple{constant_refs::Vector{Base.RefValue{ParametricNode{Float64}}}, parameter_refs::Matrix{Float64}, num_parameters::Int64, num_constants::Int64}, Dataset{Float64, Float64, Matrix{Float64}, Vector{Float64}, Nothing, @NamedTuple{class::Vector{Int64}}, Nothing, Nothing, Nothing, Nothing}, Options{SymbolicRegression.CoreModule.OptionsStructModule.ComplexityMapping{Int64, Int64}, DynamicExpressions.OperatorEnumModule.OperatorEnum{Tuple{typeof(+), typeof(*), typeof(/), typeof(-)}, Tuple{typeof(cos), typeof(exp)}}, ParametricNode, ParametricExpression, @NamedTuple{max_parameters::Int64}, MutationWeights, false, false, nothing, ADTypes.AutoEnzyme{Nothing, Nothing}, 5}, Nothing}, ADTypes.AutoEnzyme{Nothing, Nothing}, @NamedTuple{storage_tree::ParametricExpression{Float64, ParametricNode{Float64}, @NamedTuple{operators::Nothing, variable_names::Nothing, parameters::Matrix{Float64}, parameter_names::Nothing}}, storage_refs::@NamedTuple{constant_refs::Vector{Base.RefValue{ParametricNode{Float64}}}, parameter_refs::Matrix{Float64}, num_parameters::Int64, num_constants::Int64}, storage_dataset::Dataset{Float64, Float64, Matrix{Float64}, Vector{Float64}, Nothing, @NamedTuple{class::Vector{Int64}}, Nothing, Nothing, Nothing, Nothing}}}, Vector{Float64}, Vector{Float64}})()
                @ SymbolicRegressionEnzymeExt ~/PermaDocuments/SymbolicRegressionMonorepo/SymbolicRegression.jl/ext/SymbolicRegressionEnzymeExt.jl:42
in expression starting at /Users/mcranmer/PermaDocuments/SymbolicRegressionMonorepo/SymbolicRegression.jl/examples/parameterized_function.jl:101

The code is here: https://github.com/MilesCranmer/SymbolicRegression.jl/blob/667df823cd2dd111db524b7cb0c495d1a583eb89/ext/SymbolicRegressionEnzymeExt.jl#L42

[wsmoses]

Okay good news that’s now a different error

[MilesCranmer]

Same error on main. Is there any way I can get more debugging info out of this to see where abs_typeof is blowing up? As far as I know there should be no type instabilities here, and this same code used to work ok.

[wsmoses]

what's the backtrace on main (and list commit you use). The line number you have shouldn't throw that error. With that I'll try to make a patch to add more info to the relevant assertion.