Sort in augmented primal

EnzymeAD · Aug 16, 2023 · f56cbe8 · f56cbe8
1 parent 82e44ea
commit f56cbe8
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/src/internal_rules.jl b/src/internal_rules.jl
@@ -126,15 +126,17 @@ function EnzymeRules.augmented_primal(
         kwargs...
     )
     inds = sortperm(xs.val; kwargs...)
+    xs.val .= xs.val[inds]
+    xs.dval .= xs.dval[inds]
     if EnzymeRules.needs_primal(config)
-        primal = xs.val[inds]
+        primal = xs.val
     else
         primal = nothing
     end
     if RT <: Const
         shadow = nothing
     else
-        shadow = xs.dval[inds]
+        shadow = xs.dval
     end
     return EnzymeRules.AugmentedReturn(primal, shadow, inds)
 end