Type-unstable `ifelse` over tuples in `interpolating_time_indices(::Clamp, …)` (mixed `Int`/`Float64` slot)
@dkytezab I think I found it
## Summary
`interpolating_time_indices(::Clamp, times, t)` builds three index tuples whose
**first slot** (the time-interpolation fraction `ñ`) has inconsistent element
types: it is a literal `Int` `0` in the two clamped branches but a `Float64`
in the in-range branch.
```julia
beyond_indices = (0, n₂, n₂) # slot 1 :: Int
before_indices = (0, n₁, n₁) # slot 1 :: Int
unclamped_indices = (ñ, n₁, n₂) # slot 1 :: typeof(ñ) (Float)
```
These tuples are then selected with `ifelse`, so the return type is
`Union{Tuple{Int,Int,Int}, Tuple{Float64,Int,Int}}` — a small **type
instability** in ordinary execution, and under Reactant it triggers a warning
on every compile:
```
┌ Warning: `ifelse` with different element-types in Reactant works by promoting the
│ element-type to the common type. This is semantically different from the behavior
│ of `ifelse` in Base. Use with caution
└ @ Reactant.TracedRNumberOverrides .../src/TracedRNumber.jl:595
```
Because `Clamp()` is the **default** `time_indexing` for `FieldTimeSeries`, this
fires for essentially any Reactant-compiled model that interpolates a
`FieldTimeSeries` in time.
## Location
`src/OutputReaders/field_time_series_indexing.jl`, function
`interpolating_time_indices(::Clamp, times, t)` (Oceananigans v0.110.6):
```julia
# Clamp mode if out-of-bounds, i.e get the neareast neighbor
@inline function interpolating_time_indices(::Clamp, times, t)
ñ, n₁, n₂ = find_time_index(times, t)
beyond_indices = (0, n₂, n₂) # Beyond the last time: return n₂
before_indices = (0, n₁, n₁) # Before the first time: return n₁
unclamped_indices = (ñ, n₁, n₂) # Business as usual
Nt = length(times)
return ifelse(ñ + n₁ > Nt, beyond_indices,
ifelse(ñ + n₁ < 1, before_indices, unclamped_indices))
end
```
## Root cause
Reactant maps `ifelse` over tuples element-wise
(`Base.ifelse(::TracedRNumber{Bool}, ::Tuple, ::Tuple)` →
`ntuple(i -> ifelse(pred, x[i], y[i]))`). For slot 1 this becomes
`ifelse(pred, 0::Int, ñ::Float64)`, i.e. an `ifelse` whose two value branches
have different element types — which is precisely what the warning flags. The
value is correct (the `Int` `0` is promoted to `0.0`, the intended snap-to-
neighbor fraction), but the mixed types are unnecessary.
Note the sibling function already does the type-stable thing — `find_time_index`
uses `zero(ñ)` rather than `0`:
```julia
ñ = ifelse(n₂ == n₁, zero(ñ), ñ)
```
The `Cyclical` and `Linear` paths are unaffected: `Cyclical`'s tuples are all
`(Float64, Int, Int)`, and `Linear` returns the tuple directly with no
tuple-`ifelse`.
## Proposed fix
Make slot 1 type-uniform by replacing the literal `0` with `zero(ñ)`:
```diff
- beyond_indices = (0, n₂, n₂) # Beyond the last time: return n₂
- before_indices = (0, n₁, n₁) # Before the first time: return n₁
+ beyond_indices = (zero(ñ), n₂, n₂) # Beyond the last time: return n₂
+ before_indices = (zero(ñ), n₁, n₁) # Before the first time: return n₁
unclamped_indices = (ñ, n₁, n₂) # Business as usual
```
This is value-preserving (`zero(ñ) == 0`), removes the type instability in
ordinary execution, and silences the Reactant warning. I verified it both ways
(see below).
## MWE
```julia
using Reactant
using Oceananigans
using Oceananigans.OutputReaders: interpolating_time_indices, Clamp, Cyclical, Linear
using Logging
# A logger that records every message and *ignores* `maxlog`, so the one-shot
# Reactant warning (maxlog = 1) can be observed once per compile in one session.
struct CaptureLogger <: AbstractLogger
messages :: Vector{String}
end
Logging.min_enabled_level(::CaptureLogger) = Logging.BelowMinLevel
Logging.shouldlog(::CaptureLogger, args...) = true
Logging.catch_exceptions(::CaptureLogger) = false
Logging.handle_message(l::CaptureLogger, level, message, args...; kwargs...) =
(push!(l.messages, string(message)); nothing)
is_ifelse_warning(m) = occursin("ifelse", m) && occursin("different element-types", m)
function report(label, messages)
warnings = filter(is_ifelse_warning, messages)
println(" ", label, " → ifelse element-type warnings: ", length(warnings))
for (i, w) in enumerate(warnings)
println(" [$i] ", replace(w, r"\s+" => " "))
end
end
# Part 1 — the bare mechanism, no Oceananigans: ifelse over two tuples whose
# slot 1 differs in element type (Int vs Float).
mixed_tuple_ifelse(pred, x) = ifelse(pred, (0, x), (x, x))
logger1 = CaptureLogger(String[])
with_logger(logger1) do
Reactant.@compile mixed_tuple_ifelse(Reactant.ConcreteRNumber(true),
Reactant.ConcreteRNumber(2.5))
end
println("Part 1 — bare tuple ifelse (Int slot vs Float slot)")
report("mixed_tuple_ifelse", logger1.messages)
println()
# Part 2 — the real culprit: interpolating_time_indices per time_indexing mode.
times = range(0, 10, step = 1) # StepRange → closed-form searchsortedfirst (traceable)
function warnings_for(mode)
fraction(t) = first(interpolating_time_indices(mode, times, t))
logger = CaptureLogger(String[])
with_logger(logger) do
Reactant.@compile fraction(Reactant.ConcreteRNumber(2.5))
end
return logger.messages
end
println("Part 2 — Oceananigans interpolating_time_indices, by time_indexing mode")
for (name, mode) in (("Clamp (FieldTimeSeries default)", Clamp()),
("Cyclical (period 11)", Cyclical(11)),
("Linear ", Linear()))
try
report(name, warnings_for(mode))
catch err
println(" ", name, " → compile failed: ", sprint(showerror, err))
end
end
```
### Output (before the fix)
```
Part 1 — bare tuple ifelse (Int slot vs Float slot)
mixed_tuple_ifelse → ifelse element-type warnings: 1
[1] `ifelse` with different element-types in Reactant works by promoting the element-type to the common type. This is semantically different from the behavior of `ifelse` in Base. Use with caution
Part 2 — Oceananigans interpolating_time_indices, by time_indexing mode
Clamp (FieldTimeSeries default) → ifelse element-type warnings: 2
[1] `ifelse` with different element-types ...
[2] `ifelse` with different element-types ...
Cyclical (period 11) → ifelse element-type warnings: 0
Linear → ifelse element-type warnings: 0
```
### After the fix
Applying the `zero(ñ)` patch and re-running the `Clamp` case:
```
PATCHED Clamp → ifelse element-type warnings: 0
fraction(99) = 0.0 # clamped beyond last time → fraction 0, unchanged
fraction(-9) = 0.0 # clamped before first time → fraction 0, unchanged
```
## Environment
- Oceananigans v0.110.6
- Reactant v0.2.268
- Julia 1.11 / 1.12 (CPU backend)
The instability also exists without Reactant:
```julia
julia> Base.return_types(interpolating_time_indices, Tuple{Clamp, typeof(0:1:10), Float64})[1]
Union{Tuple{Float64, Int64, Int64}, Tuple{Int64, Int64, Int64}} # non-concrete
```
so the fix is a small correctness/inference improvement independent of Reactant.
0 条评论