Open wsmoses opened 11 months ago
MWE before #1207:
using Enzyme
f(a, x) = first(sortperm(a)) * x
a = [3.0, 2.0, 4.0]
da = zero(a)
autodiff(Reverse, f, Active, Duplicated(a, da), Active(5.0))
ERROR: Enzyme execution failed.
Enzyme: No custom augmented_primal rule was appliable for Tuple{Duplicated{NamedTuple{(:alg, :order, :scratch), Tuple{Base.Sort.MissingOptimization{Base.Sort.BoolOptimization{Base.Sort.Small{10, Base.Sort.InsertionSortAlg, Base.Sort.IEEEFloatOptimization{Base.Sort.IsUIntMappable{Base.Sort.Small{40, Base.Sort.InsertionSortAlg, Base.Sort.CheckSorted{Base.Sort.ComputeExtrema{Base.Sort.ConsiderCountingSort{Base.Sort.CountingSort, Base.Sort.ConsiderRadixSort{Base.Sort.RadixSort, Base.Sort.Small{80, Base.Sort.InsertionSortAlg, Base.Sort.ScratchQuickSort{Missing, Missing, Base.Sort.InsertionSortAlg}}}}}}}, Base.Sort.StableCheckSorted{Base.Sort.ScratchQuickSort{Missing, Missing, Base.Sort.InsertionSortAlg}}}}}}}, Base.Order.Perm{Base.Order.ForwardOrdering, Vector{Float64}}, Nothing}}}, typeof(EnzymeCore.EnzymeRules.augmented_primal), EnzymeCore.EnzymeRules.ConfigWidth{1, false, false, (false, true)}, Const{typeof(sort!)}, Type{Const{Vector{Int64}}}, Const{Vector{Int64}}}
Stacktrace:
[1] #_sortperm#33
@ ./sort.jl:1582
Would it be possible to define a rule for Const
activity xs
?
It would be, but we really should not write a rule in that case (since otherwise we'd unnecessarily prevent inlining/etc).
See: https://github.com/EnzymeAD/Enzyme.jl/pull/1207
@jgreener64 your sort rule was broken on integers, which resulted in issues when used, for example, on the sort! inside of a sortperm.
This PR is over conservative and changes the sort! rule to only apply to floats.
What I'd really like here is to define the sort rule only for types that are not Enzyme.guaranteed_const. Or at minimum for Vector{Int}. @vchuravy any thoughts here how?
Going to merge this regardless for now since otherwise it breaks other things (like oceananigans)