EnzymeAD / Enzyme.jl

Julia bindings for the Enzyme automatic differentiator
https://enzyme.mit.edu
MIT License
455 stars 64 forks source link

Sort rule for non-float type #1208

Open wsmoses opened 11 months ago

wsmoses commented 11 months ago

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)

jgreener64 commented 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?

wsmoses commented 10 months ago

It would be, but we really should not write a rule in that case (since otherwise we'd unnecessarily prevent inlining/etc).