Fix GPU compilation of Dense(_, _, tanh) on CuArray

[ChrisRackauckas-Claude]

Summary

Fixes the four DeepSplitting GPU tests that have been failing on the self-hosted runner since the SciMLBase v3 / CUDA 6 / NNlib 0.9.3x bump:

• DeepSplitting algorithm - allen cahn
• DeepSplitting - Black-Scholes Equation with Default Risk
• DeepSplitting - replicator mutator
• DeepSplitting algorithm - allen cahn non local - Neumann BC

(Same set that was failing on main since PR #137 / commit 6fc04330.)

Root cause

All four tests build Dense(d, hls, tanh) layers and run them on CuArray input. Inside NNlib.bias_act!, the activation is wrapped into a ComposedFunction(fast_act(σ, x), +) and broadcast on GPU. For σ = tanh, fast_act substitutes tanh_fast (a polynomial approximation built on evalpoly); the resulting ComposedFunction{tanh_fast, +} broadcast kernel then hits a dynamic dispatch on the ComposedFunction kwsorter that the CUDA compiler rejects:

InvalidIRError: compiling MethodInstance for gpu_broadcast_kernel_cartesian(...)
Reason: unsupported dynamic function invocation
  (call to var\"#_#103\"(kw::Base.Pairs{...}, c::ComposedFunction, x...)
   @ Base operators.jl:1041)

The same construction with the device intrinsic tanh (the original Base method) compiles cleanly — Carlo Lucibello reduced exactly this in FluxML/Flux.jl#2633 on Metal:

NNlib._fast_broadcast!(tanh ∘ (+), W*x, b)              # OK
NNlib._fast_broadcast!(NNlib.tanh_fast ∘ (+), W*x, b)   # InvalidIRError

Fix

NNlib already exposes a per-array-type opt-out for exactly this case via NNlib.fast_act(σ, x) (see activations.jl:897-906). Add the CuArray override in HighDimPDE so Dense(_, _, tanh) falls back to Base.tanh on the GPU. NNlib is added as a direct dep so the override is unambiguous rather than relying on Flux's transitive load.

This is the same opt-out @mcabbott proposed for Metal; the upstream Metal PR (NNlib#666) used @device_override instead but was closed unmerged. Until NNlib lands a proper CUDA-side fix, this in-package override is the smallest change that gets the GPU tests green here.

Test plan

• GPU runner picks up the new commit and DeepSplitting allen cahn, Black-Scholes Default Risk, replicator mutator, allen cahn non local - Neumann BC go from errored to passed.
• Core / Documentation continue to pass (no behaviour change on CPU because the override is gated on ::CuArray).
• No regression in CuArray DeepSplitting tests that use relu (which already passed).

🤖 Generated with Claude Code

This PR should be ignored until reviewed by @ChrisRackauckas.

[ChrisRackauckas-Claude]

Tracking down the upstream story

Metal already has the sister bug — fixed in NNlib v0.9.32

The same failure is reported on Metal as
FluxML/Flux.jl#2633 and
JuliaGPU/Metal.jl#673.
@CarloLucibello reduced it on Metal to:

using NNlib, Metal
W, x, b = Mtl.(rand(Float32, 3, 2)), Mtl.(rand(Float32, 2, 2)), Mtl.(rand(Float32, 3))

NNlib._fast_broadcast!(tanh ∘ (+),               W*x, b)  # ✅ OK
NNlib._fast_broadcast!(NNlib.tanh_fast ∘ (+),    W*x, b)  # ❌ InvalidIRError

@christiangnrd narrowed it further: removing sign(x) from a hand-rolled
copy of tanh_fast makes it compile. The fix that landed in NNlib v0.9.32
(NNlib#666) is a
single-line NNlibMetalExt:

@device_override NNlib.tanh_fast(x) = Base.FastMath.tanh_fast(x)

CUDA has the same bug — but no upstream issue or extension fix

I searched FluxML/NNlib.jl, JuliaGPU/CUDA.jl, JuliaGPU/GPUCompiler.jl,
FluxML/Flux.jl, and LuxDL/Lux.jl. Nobody has filed it for CUDA. There's
no NNlibCUDAExt tanh_fast override either — the existing CUDA ext
only carries softplus etc. activation comments and conv_bias_act!,
not the polynomial-tanh wrapper the Metal ext has.

So we are most likely the first to trip over this on CUDA. The reason
it didn't surface earlier on main is probably that the only
Dense(_, _, tanh) GPU paths in this repo are the four DeepSplitting
tests — the rest of the suite uses relu or identity.

Reduced CUDA MWE (would need GPU hardware)

using CUDA, NNlib

a = CUDA.rand(Float32, 5)
b = CUDA.rand(Float32, 5)

# ✅ OK — Base.tanh, GPU compiler uses the device intrinsic
broadcast(tanh ∘ (+), a, b)

# ❌ InvalidIRError on `var"#_#103"(kw::Base.Pairs, c::ComposedFunction, x...)`
broadcast(NNlib.tanh_fast ∘ (+), a, b)

Equivalently, what the failing tests actually trigger:

using CUDA, Flux
m = Dense(2 => 3, tanh) |> Flux.gpu
x = CUDA.rand(Float32, 2, 4)
m(x)         # forward fails with InvalidIRError

Why tanh_fast ∘ + fails but tanh ∘ + works

NNlib.tanh_fast(x::Float32) ends in
ifelse(x2 < 66f0, x*(n/d), sign(x)). The sign branch makes the body
non-trivially type-stable inside the Broadcasted{ComposedFunction{...}}
closure, and the GPU compiler then escapes to the dynamic kwsorter
(::ComposedFunction)(x...; kw...) at operators.jl:1041 instead of
inlining the call. With Base.tanh the body inlines straight to the
device tanh intrinsic, the kwsorter dispatch never has to be resolved
dynamically, and the kernel compiles.

Where the fix really belongs

Upstream NNlib: add an NNlibCUDAExt @device_override mirroring the
Metal one. Single line:

@device_override NNlib.tanh_fast(x) = Base.FastMath.tanh_fast(x)

I'll file an issue + PR against FluxML/NNlib.jl once the workaround
here is confirmed green on the GPU runner. Until that lands and ships,
the per-array-type fast_act opt-out in this PR is the smallest
in-package fix; it's exactly the shape NNlib designed fast_act's 2nd
positional arg for (activations.jl#L897-L906).

🤖 Generated with Claude Code

[ChrisRackauckas-Claude]

GPU run on 1c0769b confirmed the fast_act opt-out resolved the four InvalidIRErrors — the GPU summary went from 222 passed, 4 errored to 227 passed, 1 failed. The new failure is just Aqua's piracy check correctly flagging the override:

Possible type-piracy detected:
[1] fast_act(::typeof(tanh), ::CUDACore.CuArray) @ HighDimPDE src/HighDimPDE.jl:29

That's intentional — we're adding a method to NNlib's fast_act for ::CuArray, which is exactly the per-array-type opt-out shape NNlib designed fast_act's 2nd positional arg for. Pushed 712fea1 whitelisting it via Aqua.test_piracies(HighDimPDE, treat_as_own = [NNlib.fast_act]). Once that goes green I'll file the upstream NNlib issue + PR proposing the equivalent NNlibCUDAExt @device_override.