Add Reactant extension for branchless interpolation#516
Add Reactant extension for branchless interpolation#516ChrisRackauckas-Claude wants to merge 1 commit into
Conversation
DataInterpolations methods use `if t < first(A.t)` branching that fails when
`t` is a Reactant.TracedRNumber (produces TracedRNumber{Bool} which can't be
used in boolean context). This adds a package extension that provides branchless
`_interpolate` methods using `ifelse` chains for LinearInterpolation,
ConstantInterpolation, and QuadraticInterpolation.
Fixes SciML/SciMLSensitivity.jl#1430
Co-Authored-By: Chris Rackauckas <accounts@chrisrackauckas.com>
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Investigation: alternative approachesExplored several alternatives to the branchless
|
| Knots | Branching (ns) | Branchless (ns) | Ratio |
|---|---|---|---|
| 5 | 54 | 19 | 0.35x |
| 10 | 53 | 32 | 0.61x |
| 20 | 53 | 58 | ~1.1x |
| 50 | 53 | 143 | 2.7x |
| 100 | 54 | 270 | 5x |
| 1000 | 54 | 2827 | 52x |
The branching version is O(log n) via correlated binary search. The branchless O(n) scan wins for small tables but degrades linearly. For typical ODE interpolation tables (<100 knots), the overhead is acceptable.
Conclusion
The branchless TracedRNumber dispatch approach is the only viable option from an extension. The fundamental issue is that Reactant's call_with_reactant generated function traces through function bodies directly, and no mechanism (@trace, overlay, traced_call) can prevent it from encountering the if statement in _interpolate. The only way to avoid this is to dispatch to an entirely different method body that never uses if on traced values.
A deeper fix would require either:
- DataInterpolations adding
ReactantCoreas a dependency and sprinkling@traceonifstatements in source (but@trace ifneeds both branches to be traceable, ruling outthrowbranches) - Reactant learning to automatically convert
ifwithTracedRNumber{Bool}conditions totraced_ifduring its overlay pass (a Reactant.jl enhancement)
|
@avik-pal I'm surprised an overlay doesn't end up good here? |
Summary
DataInterpolationsReactantExtpackage extension that provides branchless_interpolatemethods forTracedRNumberinputsLinearInterpolation,ConstantInterpolation, andQuadraticInterpolationifelsechains (lowered tostablehlo.select) instead ofifbranching, which fails withTracedRNumber{Bool}ExtrapolationTypevariants (Constant, Linear, Extension) in the traced pathMotivation
Fixes SciML/SciMLSensitivity.jl#1430 — when using
ReactantVJPas theautojacvecinGaussAdjoint, DataInterpolations calls inside ODE functions fail because_interpolateusesif t < first(A.t)wheretis aReactant.TracedRNumber. The comparison producesTracedRNumber{Bool}which can't be used in Julia's boolean context.Approach
For each supported interpolation type, the extension evaluates all segments and selects the correct one using
ifelse. This is O(n) in knot count but fully traceable by Reactant. For typical ODE interpolation tables (tens to hundreds of knots), this is efficient enough.Test plan
🤖 Generated with Claude Code