Floating point error using adjoint_sensitivities when tspan includes 16f0

[jesseylin]

Describe the bug 🐞

Adjoint solve immediately fails with a floating point precision error exactly at 16f0, regardless of the ODEFunction, solver, etc.

Expected behavior

The solve should not fail, or at least its propensity to fail is smooth in tspan

Minimal Reproducible Example 👇

using Lux
using Random
using Zygote
using OrdinaryDiffEq
using SciMLSensitivity
using ComponentArrays
using Optimisers
using LinearAlgebra

using Test

# extremely basic test of SciMLSensitivities
rng = Xoshiro(1234)

function flow(model_spec; tspan=(0.0f0, 20.0f0))
    return @compact(; model_spec=model_spec) do s, ps
        dudt(u, p, t) = -u + model_spec(reshape(u, size(s)), p)
        prob =
            ODEProblem(ODEFunction{false}(dudt), vec(s), tspan, ps.model_spec)
        @return solve(prob, Tsit5())
    end
end

model_spec = Dense(11 => 11)
model = flow(model_spec; tspan=(0.0f0, prevfloat(16.0f0)))
_ps, st = Lux.setup(rng, model)
ps = ComponentArray(_ps)
s0 = randn(rng, 11)
sol, st = model(s0, ps, st)
sensealg = GaussAdjoint(autojacvec=ReverseDiffVJP())

# this one works
_, gs = adjoint_sensitivities(sol, Tsit5(); sensealg, g=(u, p, t) -> norm(u))
@test !iszero(gs)

# this one doesn't
model = flow(model_spec; tspan=(0.0f0, 16.0f0))
sol, st = model(s0, ps, st)
_, gs = adjoint_sensitivities(sol, Tsit5(); sensealg, g=(u, p, t) -> norm(u))
@test !iszero(gs)

Error & Stacktrace ⚠️
In the above MRE, a warning is thrown and the gradients are identically zero due to the adjoint solve failing.

┌ Warning: At t=16.0, dt was forced below floating point epsilon -1.0e-6, and step error estimate = 1.0. Aborting. There is either an error in your model specification or the true solution is unstable (or the true solution can not be represented in the precision of Float64).
└ @ SciMLBase ~/.julia/packages/SciMLBase/tEuIM/src/integrator_interface.jl:609
Test Failed at /Users/jesselin/src/mwe/mwe.jl:40
  Expression: !(iszero(gs))

ERROR: LoadError: There was an error during testing
in expression starting at /Users/jesselin/src/mwe/mwe.jl:40

Environment (please complete the following information):

• Output of using Pkg; Pkg.status()

 [6e4b80f9] BenchmarkTools v1.5.0
⌃ [13f3f980] CairoMakie v0.12.9
⌃ [f68482b8] Cthulhu v2.15.0
⌃ [a93c6f00] DataFrames v1.6.1
⌃ [634d3b9d] DrWatson v2.13.0
⌃ [5903a43b] Infiltrator v1.6.4
⌃ [98e50ef6] JuliaFormatter v1.0.60
⌃ [2b0e0bc5] LanguageServer v4.5.0
  [5fb14364] OhMyREPL v0.5.28
⌃ [e4faabce] PProf v3.1.2
⌃ [295af30f] Revise v3.5.13
⌃ [1986cc42] Unitful v1.20.0
  [9abbd945] Profile
  [8dfed614] Test
Info Packages marked with ⌃ have new versions available and may be upgradable.

• Output of using Pkg; Pkg.status(; mode = PKGMODE_MANIFEST)

Status `~/.julia/environments/v1.10/Manifest.toml`
  [621f4979] AbstractFFTs v1.5.0
  [1520ce14] AbstractTrees v0.4.5
  [79e6a3ab] Adapt v4.0.4
⌃ [35492f91] AdaptivePredicates v1.1.1
  [66dad0bd] AliasTables v1.1.3
  [27a7e980] Animations v0.4.1
  [67c07d97] Automa v1.0.4
  [13072b0f] AxisAlgorithms v1.1.0
  [39de3d68] AxisArrays v0.4.7
  [6e4b80f9] BenchmarkTools v1.5.0
  [e1450e63] BufferedStreams v1.2.2
  [fa961155] CEnum v0.5.0
  [00ebfdb7] CSTParser v3.4.3
  [159f3aea] Cairo v1.1.0
⌃ [13f3f980] CairoMakie v0.12.9
⌃ [d360d2e6] ChainRulesCore v1.24.0
  [da1fd8a2] CodeTracking v1.3.6
  [944b1d66] CodecZlib v0.7.6
  [a2cac450] ColorBrewer v0.4.0
  [35d6a980] ColorSchemes v3.26.0
  [3da002f7] ColorTypes v0.11.5
  [c3611d14] ColorVectorSpace v0.10.0
  [5ae59095] Colors v0.12.11
⌃ [a80b9123] CommonMark v0.8.12
  [34da2185] Compat v4.16.0
  [187b0558] ConstructionBase v1.5.8
  [d38c429a] Contour v0.6.3
  [a8cc5b0e] Crayons v4.1.1
⌃ [f68482b8] Cthulhu v2.15.0
  [9a962f9c] DataAPI v1.16.0
⌃ [a93c6f00] DataFrames v1.6.1
  [864edb3b] DataStructures v0.18.20
  [e2d170a0] DataValueInterfaces v1.0.0
⌃ [927a84f5] DelaunayTriangulation v1.3.0
⌃ [31c24e10] Distributions v0.25.111
  [ffbed154] DocStringExtensions v0.9.3
⌃ [634d3b9d] DrWatson v2.13.0
  [4e289a0a] EnumX v1.0.4
  [429591f6] ExactPredicates v2.2.8
  [411431e0] Extents v0.1.4
  [7a1cc6ca] FFTW v1.8.0
  [5789e2e9] FileIO v1.16.3
  [8fc22ac5] FilePaths v0.8.3
  [48062228] FilePathsBase v0.9.22
  [1a297f60] FillArrays v1.13.0
  [53c48c17] FixedPointNumbers v0.8.5
  [08572546] FlameGraphs v1.0.0
  [1eca21be] FoldingTrees v1.2.1
  [1fa38f19] Format v1.3.7
  [b38be410] FreeType v4.1.1
  [663a7486] FreeTypeAbstraction v0.10.3
  [68eda718] GeoFormatTypes v0.4.2
⌃ [cf35fbd7] GeoInterface v1.3.6
  [5c1252a2] GeometryBasics v0.4.11
  [c27321d9] Glob v1.3.1
  [a2bd30eb] Graphics v1.1.2
  [3955a311] GridLayoutBase v0.11.0
  [42e2da0e] Grisu v1.0.2
  [34004b35] HypergeometricFunctions v0.3.24
  [2803e5a7] ImageAxes v0.6.11
  [c817782e] ImageBase v0.1.7
  [a09fc81d] ImageCore v0.10.2
  [82e4d734] ImageIO v0.6.8
  [bc367c6b] ImageMetadata v0.9.9
  [9b13fd28] IndirectArrays v1.0.0
⌃ [5903a43b] Infiltrator v1.6.4
  [d25df0c9] Inflate v0.1.5
  [842dd82b] InlineStrings v1.4.2
  [a98d9a8b] Interpolations v0.15.1
⌃ [d1acc4aa] IntervalArithmetic v0.22.16
  [8197267c] IntervalSets v0.7.10
  [41ab1584] InvertedIndices v1.3.0
  [92d709cd] IrrationalConstants v0.2.2
  [f1662d9f] Isoband v0.1.1
  [c8e1da08] IterTools v1.10.0
  [82899510] IteratorInterfaceExtensions v1.0.0
⌅ [033835bb] JLD2 v0.4.53
  [1019f520] JLFzf v0.1.8
  [692b3bcd] JLLWrappers v1.6.0
  [682c06a0] JSON v0.21.4
  [b9b8584e] JSONRPC v1.4.1
  [b835a17e] JpegTurbo v0.1.5
⌃ [98e50ef6] JuliaFormatter v1.0.60
⌃ [aa1ae85d] JuliaInterpreter v0.9.35
  [70703baa] JuliaSyntax v0.4.10
  [5ab0869b] KernelDensity v0.6.9
  [b964fa9f] LaTeXStrings v1.3.1
⌃ [2b0e0bc5] LanguageServer v4.5.0
  [8cdb02fc] LazyModules v0.3.1
  [1d6d02ad] LeftChildRightSiblingTrees v0.2.0
  [2ab3a3ac] LogExpFunctions v0.3.28
⌅ [6f1432cf] LoweredCodeUtils v2.4.8
  [1914dd2f] MacroTools v0.5.13
⌅ [ee78f7c6] Makie v0.21.9
⌅ [20f20a25] MakieCore v0.8.6
  [dbb5928d] MappedArrays v0.4.2
  [0a4f8689] MathTeXEngine v0.6.1
  [e1d29d7a] Missings v1.2.0
  [e94cdb99] MosaicViews v0.3.4
  [77ba4419] NaNMath v1.0.2
  [f09324ee] Netpbm v1.1.1
  [510215fc] Observables v0.5.5
  [6fe1bfb0] OffsetArrays v1.14.1
  [5fb14364] OhMyREPL v0.5.28
  [52e1d378] OpenEXR v0.3.2
  [bac558e1] OrderedCollections v1.6.3
  [90014a1f] PDMats v0.11.31
  [f57f5aa1] PNGFiles v0.4.3
⌃ [e4faabce] PProf v3.1.2
  [19eb6ba3] Packing v0.5.0
  [5432bcbf] PaddedViews v0.5.12
  [69de0a69] Parsers v2.8.1
  [b98c9c47] Pipe v1.3.0
  [eebad327] PkgVersion v0.3.3
  [995b91a9] PlotUtils v1.4.1
  [647866c9] PolygonOps v0.1.2
  [2dfb63ee] PooledArrays v1.4.3
  [aea7be01] PrecompileTools v1.2.1
  [21216c6a] Preferences v1.4.3
⌃ [08abe8d2] PrettyTables v2.3.2
  [92933f4c] ProgressMeter v1.10.2
  [3349acd9] ProtoBuf v1.0.16
  [43287f4e] PtrArrays v1.2.1
  [4b34888f] QOI v1.0.0
⌃ [1fd47b50] QuadGK v2.11.0
  [b3c3ace0] RangeArrays v0.3.2
  [c84ed2f1] Ratios v0.4.5
  [189a3867] Reexport v1.2.2
  [05181044] RelocatableFolders v1.0.1
  [ae029012] Requires v1.3.0
⌃ [295af30f] Revise v3.5.13
⌅ [79098fc4] Rmath v0.7.1
  [5eaf0fd0] RoundingEmulator v0.2.1
⌃ [fdea26ae] SIMD v3.5.0
  [6c6a2e73] Scratch v1.2.1
  [91c51154] SentinelArrays v1.4.5
  [65257c39] ShaderAbstractions v0.4.1
  [992d4aef] Showoff v1.0.3
  [73760f76] SignedDistanceFields v0.4.0
  [699a6c99] SimpleTraits v0.9.4
  [45858cf5] Sixel v0.1.3
  [a2af1166] SortingAlgorithms v1.2.1
  [276daf66] SpecialFunctions v2.4.0
  [cae243ae] StackViews v0.1.1
  [90137ffa] StaticArrays v1.9.7
  [1e83bf80] StaticArraysCore v1.4.3
  [b3cc710f] StaticLint v8.2.2
  [82ae8749] StatsAPI v1.7.0
⌃ [2913bbd2] StatsBase v0.33.21
⌃ [4c63d2b9] StatsFuns v1.3.1
⌅ [892a3eda] StringManipulation v0.3.4
  [09ab397b] StructArrays v0.6.18
⌅ [cf896787] SymbolServer v7.4.0
  [3783bdb8] TableTraits v1.0.1
  [bd369af6] Tables v1.12.0
  [62fd8b95] TensorCore v0.1.1
⌅ [76b0de8b] TestItemDetection v0.2.0
  [731e570b] TiffImages v0.10.0
  [0796e94c] Tokenize v0.5.29
⌃ [3bb67fe8] TranscodingStreams v0.11.2
  [981d1d27] TriplotBase v0.1.0
⌃ [d265eb64] TypedSyntax v1.4.1
  [5c2747f8] URIs v1.5.1
  [3a884ed6] UnPack v1.0.2
  [1cfade01] UnicodeFun v0.4.1
⌃ [1986cc42] Unitful v1.20.0
  [b8c1c048] WidthLimitedIO v1.0.1
  [efce3f68] WoodburyMatrices v1.0.0
  [6e34b625] Bzip2_jll v1.0.8+1
  [4e9b3aee] CRlibm_jll v1.0.1+0
⌃ [83423d85] Cairo_jll v1.18.0+2
  [5ae413db] EarCut_jll v2.2.4+0
  [2e619515] Expat_jll v2.6.2+0
  [b22a6f82] FFMPEG_jll v6.1.2+0
⌃ [f5851436] FFTW_jll v3.3.10+0
  [a3f928ae] Fontconfig_jll v2.13.96+0
  [d7e528f0] FreeType2_jll v2.13.2+0
  [559328eb] FriBidi_jll v1.0.14+0
  [78b55507] Gettext_jll v0.21.0+0
⌃ [7746bdde] Glib_jll v2.80.2+0
  [3b182d85] Graphite2_jll v1.3.14+0
  [3c863552] Graphviz_jll v2.50.0+1
  [2e76f6c2] HarfBuzz_jll v8.3.1+0
  [905a6f67] Imath_jll v3.1.11+0
  [1d5cc7b8] IntelOpenMP_jll v2024.2.1+0
⌃ [aacddb02] JpegTurbo_jll v3.0.3+0
  [c1c5ebd0] LAME_jll v3.100.2+0
⌃ [1d63c593] LLVMOpenMP_jll v17.0.6+0
⌃ [dd4b983a] LZO_jll v2.10.2+0
⌅ [e9f186c6] Libffi_jll v3.2.2+1
  [d4300ac3] Libgcrypt_jll v1.8.11+0
  [7add5ba3] Libgpg_error_jll v1.49.0+0
  [94ce4f54] Libiconv_jll v1.17.0+0
  [4b2f31a3] Libmount_jll v2.40.1+0
  [38a345b3] Libuuid_jll v2.40.1+0
  [856f044c] MKL_jll v2024.2.0+0
  [e7412a2a] Ogg_jll v1.3.5+1
  [18a262bb] OpenEXR_jll v3.2.4+0
⌃ [458c3c95] OpenSSL_jll v3.0.15+0
  [efe28fd5] OpenSpecFun_jll v0.5.5+0
  [91d4177d] Opus_jll v1.3.3+0
  [36c8627f] Pango_jll v1.54.1+0
  [30392449] Pixman_jll v0.43.4+0
⌅ [f50d1b31] Rmath_jll v0.4.3+0
  [02c8fc9c] XML2_jll v2.13.3+0
  [aed1982a] XSLT_jll v1.1.41+0
  [4f6342f7] Xorg_libX11_jll v1.8.6+0
  [0c0b7dd1] Xorg_libXau_jll v1.0.11+0
  [a3789734] Xorg_libXdmcp_jll v1.1.4+0
  [1082639a] Xorg_libXext_jll v1.3.6+0
  [ea2f1a96] Xorg_libXrender_jll v0.9.11+0
  [14d82f49] Xorg_libpthread_stubs_jll v0.1.1+0
  [c7cfdc94] Xorg_libxcb_jll v1.17.0+0
  [c5fb5394] Xorg_xtrans_jll v1.5.0+0
  [214eeab7] fzf_jll v0.53.0+0
  [9a68df92] isoband_jll v0.2.3+0
  [a4ae2306] libaom_jll v3.9.0+0
  [0ac62f75] libass_jll v0.15.2+0
  [f638f0a6] libfdk_aac_jll v2.0.3+0
⌃ [b53b4c65] libpng_jll v1.6.43+1
⌃ [075b6546] libsixel_jll v1.10.3+0
  [f27f6e37] libvorbis_jll v1.3.7+2
  [1317d2d5] oneTBB_jll v2021.12.0+0
  [cf2c5f97] pprof_jll v1.0.1+0
  [1270edf5] x264_jll v10164.0.0+0
⌅ [dfaa095f] x265_jll v3.6.0+0
  [0dad84c5] ArgTools v1.1.1
  [56f22d72] Artifacts
  [2a0f44e3] Base64
  [8bf52ea8] CRC32c
  [ade2ca70] Dates
  [8ba89e20] Distributed
  [f43a241f] Downloads v1.6.0
  [7b1f6079] FileWatching
  [9fa8497b] Future
  [b77e0a4c] InteractiveUtils
  [4af54fe1] LazyArtifacts
  [b27032c2] LibCURL v0.6.4
  [76f85450] LibGit2
  [8f399da3] Libdl
  [37e2e46d] LinearAlgebra
  [56ddb016] Logging
  [d6f4376e] Markdown
  [a63ad114] Mmap
  [ca575930] NetworkOptions v1.2.0
  [44cfe95a] Pkg v1.10.0
  [de0858da] Printf
  [9abbd945] Profile
  [3fa0cd96] REPL
  [9a3f8284] Random
  [ea8e919c] SHA v0.7.0
  [9e88b42a] Serialization
  [1a1011a3] SharedArrays
  [6462fe0b] Sockets
  [2f01184e] SparseArrays v1.10.0
  [10745b16] Statistics v1.10.0
  [4607b0f0] SuiteSparse
  [fa267f1f] TOML v1.0.3
  [a4e569a6] Tar v1.10.0
  [8dfed614] Test
  [cf7118a7] UUIDs
  [4ec0a83e] Unicode
  [e66e0078] CompilerSupportLibraries_jll v1.1.1+0
  [deac9b47] LibCURL_jll v8.4.0+0
  [e37daf67] LibGit2_jll v1.6.4+0
  [29816b5a] LibSSH2_jll v1.11.0+1
  [c8ffd9c3] MbedTLS_jll v2.28.2+1
  [14a3606d] MozillaCACerts_jll v2023.1.10
  [4536629a] OpenBLAS_jll v0.3.23+4
  [05823500] OpenLibm_jll v0.8.1+2
  [efcefdf7] PCRE2_jll v10.42.0+1
  [bea87d4a] SuiteSparse_jll v7.2.1+1
  [83775a58] Zlib_jll v1.2.13+1
  [8e850b90] libblastrampoline_jll v5.8.0+1
  [8e850ede] nghttp2_jll v1.52.0+1
  [3f19e933] p7zip_jll v17.4.0+2
Info Packages marked with ⌃ and ⌅ have new versions available. Those with ⌃ may be upgradable, but those with ⌅ are restricted by compatibility constraints from upgrading. To see why use `status --outdated -m`

• Output of versioninfo()

Julia Version 1.10.4
Commit 48d4fd48430 (2024-06-04 10:41 UTC)
Build Info:
  Official https://julialang.org/ release
Platform Info:
  OS: macOS (arm64-apple-darwin22.4.0)
  CPU: 8 × Apple M1 Pro
  WORD_SIZE: 64
  LIBM: libopenlibm
  LLVM: libLLVM-15.0.7 (ORCJIT, apple-m1)
Threads: 6 default, 0 interactive, 3 GC (on 6 virtual cores)
Environment:
  JULIA_NUM_THREADS = auto

Additional context

Add any other context about the problem here.