Type instability in reduction operators #4869
Description
It seems that reduction operators like Average and Integral are not type stable. This is a problem both for autodifferentiability (Enzyme chokes when using Average to compute the mean of a field, for example) and for processes that rely on integrated quantities; e.g. some land surface processes rely on balance equations with terms that involve integrating over the ground domain.
I haven't looked very deeply into the code yet, so I can't speculate as to the cause.
MWE:
using Oceananigans
using Statistics
import Random
const prng = Random.MersenneTwister(1234)
grid = RectilinearGrid(CPU(), size=(1,10), x=(0,1), z=(0,1), topology=(Periodic, Flat, Bounded))
X = CenterField(grid)
X .= randn(prng, size(X))
function mean_square_field(X::Field)
Y = X^2
return Field(Average(Y, dims=(1,2,3)))[1,1,1]
end
function mean_square_field!(Ȳ::Field, X::Field)
Y = X^2
set!(Ȳ, Field(Average(Y, dims=(1,2,3))))
return nothing
end
@assert mean_square_field(X) ≈ mean(X^2)
@code_warntype mean_square_field(X)
Ȳ = CenterField(grid)
@code_warntype mean_square_field!(Ȳ, X)
function integrate_field(Ȳ::Field, X::Field)
set!(Ȳ, Field(Integral(X, dims=(1,2,3))))
return nothing
end
@code_warntype integrate_field(Ȳ, X)Verified on Oceananigans v0.100.6
julia> versioninfo()
Julia Version 1.10.10
Commit 95f30e51f41 (2025-06-27 09:51 UTC)
Build Info:
Official https://julialang.org/ release
Platform Info:
OS: Linux (x86_64-linux-gnu)
CPU: 128 × AMD EPYC 9354 32-Core Processor
WORD_SIZE: 64
LIBM: libopenlibm
LLVM: libLLVM-15.0.7 (ORCJIT, znver3)
Threads: 1 default, 0 interactive, 1 GC (on 128 virtual cores)
Environment:
JULIA_NUM_THREADS =
JULIA_VSCODE_REPL = 1
JULIA_EDITOR = code
Possibly related to #4633