ITADN

Type-unstable `ifelse` over tuples in `interpolating_time_indices(::Clamp, …)` (mixed `Int`/`Float64` slot)

#5731Openxkykai 创建于 2026-06-27
X
xkykaicommented
@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 条评论