DifferentialEquations Package Kills Performance Everywhere when TruncatedStacktraces is used

The Nothing parameters are used when no value was specified. For example one of those Nothings exists because there are no tstops (times where you want to force the solver to stop). Sure you could express this as an empty array, but if you did so, you still wouldn’t know the type of the array it should be so you would still need a type parameter.

I understand that. This reminds me of the admonition in the Elm Lang book to not overuse the Maybe type:

This Maybe type is quite useful, but there are limits. Beginners are particularly prone to getting excited about Maybe and using it everywhere, even though a custom type would be more appropriate.

In other words, types that are 90% empty are not doing a very good job of encoding the domain logic into the type system. Julia has an expressive type system. Filling 90% of fields with nothing is not really taking advantage of the expressiveness of the type system.

2 Likes

How so?
You don’t want to define structs representing the exponential explosion of possibilities. Instead, use a large product type parameterized to express whatever particular combinations of behavior or functionality the user wants.

1 Like

Perhaps a less disruptive solution to address the stacktrace readability might be exported type-aliases for common cases? E.g.

julia> module temp
           struct A{T1,T2,T3}
               a::T1
               b::T2
               c::T3
           end
           AN{T1} = A{T1,Nothing,Nothing}
           export AN
       end
Main.temp

julia> f(x) = temp.A(x,nothing,nothing) + 1
f (generic function with 2 methods)

julia> f(2)
ERROR: MethodError: no method matching +(::Main.temp.AN{Int64}, ::Int64)

Closest candidates are:
  +(::Any, ::Any, ::Any, ::Any...)
   @ Base operators.jl:578
  +(::T, ::T) where T<:Union{Int128, Int16, Int32, Int64, Int8, UInt128, UInt16, UInt32, UInt64, UInt8}
   @ Base int.jl:87
  +(::Rational, ::Integer)
   @ Base rational.jl:327
  ...

Stacktrace:
 [1] f(x::Int64)
   @ Main ./REPL[17]:1
 [2] top-level scope
   @ REPL[18]:1

In this case, the alias is used in the stacktrace, and it’s much easier to read.

3 Likes

We’re discussing two missing features in the language. The problem is that type parameters for specialization are not necessarily type parameters that are intended to be used for dispatch.

So yes it’s weird that these are shown to the user at all, because there’s no intention for anyone to be dispatching on them except for the compiler as a code generation tool. However, Julia does not make a distinction between type parameters for specialization and dispatch, and so since we have not made the distinction the only way to make this optimize is to expose type parameters that have no intended use.

So then the next missing feature is control of stacktrace printing. The expected controls were for some reason not merged even as experimental:

And so we both have a bunch of type parameters that we don’t want, and we cannot control the printing to remove these pseudo type parameters.

Note that TruncatedStacktraces is now not on by default and requires a preference to be set.

## Enabling TruncatedStacktraces.jl

TruncatedStacktraces.jl is currently disabled by default, as it causes invalidations which will slow down package loading.

It can be enabled using Preferences.jl. To enable it, create a `LocalPreferences.toml` with the following entry:

[TruncatedStacktraces] disable = false

Alternatively, you can generate the `LocalPreferences.toml` using:

using Preferences, UUIDs using TruncatedStacktraces Preferences.set_preferences!(TruncatedStacktraces, "disable" => false) # OR if you don't want to load TruncatedStacktraces.jl Preferences.set_preferences!(UUID("781d530d-4396-4725-bb49-402e4bee1e77"), "disable" => false)

In either case, you need to reload your packages (depending on TruncatedStacktraces) for the change to take effect.

**TruncatedStacktraces is known to create invalidations, to remove these simply set the preference to disable it!**

It’s fine to use with system images (which everyone should be using anyways) but no longer the default. There just is no answer to better stacktraces until Julia Base merges something.

As it’s no longer the default, the loading performance issue is gone unless someone opts-in via a preference, so we can call this closed.

You do realize that in just that case you’re telling me that creating 1.3076744e+12 types on just the ODEFunction is a better way of using Julia’s expressive type system? To me, 10^12 types sounds bonkers :sweat_smile:. And note that will overflow the limit in the compiler so it’s not feasible without a custom Julia build that bumps the type limit. I would hardly say that’s a good ergonomic solution for what’s already one parameterized type. This has to be the oddest suggestion I’ve seen on the forum.

The real thing is that a form of OpaqueClosures with standard optimizations is required in Base to reduce function specializations, but that of course is a topic we all have beaten to death elsewhere.

4 Likes

Hopefully we can eventually get RFC: "quiet" types for easier-to-read stack traces by JeffBezanson · Pull Request #48444 · JuliaLang/julia · GitHub. Just looking through the PR it’s not clear to me what the drawback from including it would be, and it’s really helpful - would be nice to just add it and then think about more sophisticated ways about handling stacktrace issues further along (which I don’t imagine will be solved quite as easily or as soon especially). Am I missing something with this PR?

The daunting stacktraces are too problematic for newcomers I reckon, as I’ve seen when teaching with Julia, it’s a shame to not have #48444.

This is a good observation. Not only are long stacktraces visually noisy; to the uninitiated, they’re downright scary. Oh no! Did I break Julia!? When you don’t know what a stacktrace is good for, it’s just the computer screaming at you.

This might be heresy, but considering that 90% of the time in interactive mode I don’t care to see the stacktrace anyway because I have some silly typo that can be fixed in half a second, I’d be happy if the stacktrace didn’t even show by default. Having to type err to see the stacktrace isn’t so bad—if I’m looking at the stacktrace, I’m signing up for serious debug effort anyway.

11 Likes

Creating a new type for every possible combination of input types was not what I had in mind. What I had in mind was more composition of smaller types, most of which are concrete.

Are there fields in DEOptions or ODEIntegrator that could be a Union type instead of a type parameter? E.g., Union{Int, Nothing} or Union{Vector{Float64}, Nothing}?

I tried out the function barrier approach that I mentioned above, and it almost works, but not quite. As expected, all of the types inside the function barrier are fully inferred. But it doesn’t help with stacktraces. It just shifts all the types from one place to another.

The toy example I created is in the spirit of ODEFunction. I tried two different versions, foo and bar.

struct Funcs{F}
    f::F
    more_funcs
end

Funcs(f; more_funcs...) = Funcs(f, more_funcs)

foo(fs::Funcs, x) = _foo(fs.f, x; fs.more_funcs...)

function _foo(f, x; more_funcs...)
    g = get(more_funcs, :g, identity)
    h = get(more_funcs, :h, identity)
    k = get(more_funcs, :k, identity)
    f(g(h(k(x))))
end

function bar(fs::Funcs, x)
    g = get(fs.more_funcs, :g, identity)
    h = get(fs.more_funcs, :h, identity)
    k = get(fs.more_funcs, :k, identity)
    _bar(fs.f, g, h, k, x)
end

_bar(f, g, h, k, x) = f(g(h(k(x))))

funcs = Funcs(x -> x + 1,
    g = x -> x + 2,
    h = x -> x + 3,
    k = x -> x + 4,
)

Type inference works inside _foo and _bar:

Output of `@code_warntype`
julia> @code_warntype _foo(funcs.f, 10; funcs.more_funcs...)
MethodInstance for (::var"#_foo##kw")(::NamedTuple{(:g, :h, :k), Tuple{var"#8#12", var"#9#13", var"#10#14"}}, ::typeof(_foo), ::var"#7#11", ::Int64)
  from (::var"#_foo##kw")(::Any, ::typeof(_foo), f, x) in Main at /Users/cameron/projects/julia_misc/function_barrier.jl:11
Arguments
  _::Core.Const(var"#_foo##kw"())
  @_2::Core.Const((g = var"#8#12"(), h = var"#9#13"(), k = var"#10#14"()))
  @_3::Core.Const(_foo)
  f::Core.Const(var"#7#11"())
  x::Int64
Locals
  more_funcs...::Base.Pairs{Symbol, Function, Tuple{Symbol, Symbol, Symbol}, NamedTuple{(:g, :h, :k), Tuple{var"#8#12", var"#9#13", var"#10#14"}}}
Body::Int64
1 ─      (more_funcs... = Base.pairs(@_2))
│   %2 = Main.:(var"#_foo#6")(more_funcs...::Core.Const(Base.Pairs{Symbol, Function, Tuple{Symbol, Symbol, Symbol}, NamedTuple{(:g, :h, :k), Tuple{var"#8#12", var"#9#13", var"#10#14"}}}(:g => var"#8#12"(), :h => var"#9#13"(), :k => var"#10#14"())), @_3, f, x)::Int64
└──      return %2


julia> @code_warntype _bar(funcs.f, funcs.more_funcs[:g], funcs.more_funcs[:h], funcs.more_funcs[:k], 10)
MethodInstance for _bar(::var"#7#11", ::var"#8#12", ::var"#9#13", ::var"#10#14", ::Int64)
  from _bar(f, g, h, k, x) in Main at /Users/cameron/projects/julia_misc/function_barrier.jl:25
Arguments
  #self#::Core.Const(_bar)
  f::Core.Const(var"#7#11"())
  g::Core.Const(var"#8#12"())
  h::Core.Const(var"#9#13"())
  k::Core.Const(var"#10#14"())
  x::Int64
Body::Int64
1 ─ %1 = (k)(x)::Int64
│   %2 = (h)(%1)::Int64
│   %3 = (g)(%2)::Int64
│   %4 = (f)(%3)::Int64
└──      return %4

But unfortunately we still see all the types that we don’t want to see in the stacktrace:

foo stacktrace:

julia> foo(funcs, "a")
ERROR: MethodError: no method matching +(::String, ::Int64)
Stacktrace:
 [1] (::var"#10#14")(x::String)
   @ Main ~/projects/julia_misc/function_barrier.jl:30
 [2] _foo(f::var"#7#11", x::String; more_funcs::Base.Pairs{Symbol, Function, Tuple{Symbol, Symbol, Symbol}, NamedTuple{(:g, :h, :k), Tuple{var"#8#12", var"#9#13", var"#10#14"}}})
   @ Main ~/projects/julia_misc/function_barrier.jl:15
 [3] foo(fs::Funcs{var"#7#11"}, x::String)
   @ Main ~/projects/julia_misc/function_barrier.jl:9
 [4] top-level scope
   @ REPL[8]:1

bar stacktrace:

julia> bar(funcs, "a")
ERROR: MethodError: no method matching +(::String, ::Int64)
Stacktrace:
 [1] (::var"#10#14")(x::String)
   @ Main ~/projects/julia_misc/function_barrier.jl:30
 [2] _bar(f::var"#7#11", g::var"#8#12", h::var"#9#13", k::var"#10#14", x::String)
   @ Main ~/projects/julia_misc/function_barrier.jl:25
 [3] bar(fs::Funcs{var"#7#11"}, x::String)
   @ Main ~/projects/julia_misc/function_barrier.jl:22
 [4] top-level scope
   @ REPL[9]:1

In fact, foo is actually worse because it adds all the extra type parameters of Base.Pairs.

2 Likes

I think I figured out how to improve the printing of ODEFunction in stacktraces. The fact that the number of non-nothing fields in ODEFunction is highly variable indicates that what you really need is a collection, not a struct. So I’ve created an example type called MyFunction that wraps a named tuple. The named tuple contains only the options that are actually passed to the MyFunction constructor. For comparison, I also define a SciMLFunction type that is essentially the same as ODEFunction.

Here is the implementation for SciMLFunction and MyFunction:

struct SciMLFunction{iip, specialize, F, TMM, Ta, Tt,
                   TJ, JVP, VJP, JP, SP, TW, TWt,
                   TPJ, S, S2, S3, O, TCV, SYS}
    f::F
    mass_matrix::TMM
    analytic::Ta
    tgrad::Tt
    jac::TJ
    jvp::JVP
    vjp::VJP
    jac_prototype::JP
    sparsity::SP
    Wfact::TW
    Wfact_t::TWt
    paramjac::TPJ
    syms::S
    indepsym::S2
    paramsyms::S3
    observed::O
    colorvec::TCV
    sys::SYS
end

function SciMLFunction(f; mass_matrix=nothing, analytic=nothing, tgrad=nothing,
                       jac=nothing, jvp=nothing, vjp=nothing, jac_prototype=nothing,
                       sparsity=nothing, Wfact=nothing, Wfact_t=nothing,
                       paramjac=nothing, syms=nothing, indepsym=nothing,
                       paramsyms=nothing, observed=nothing, colorvec=nothing,
                       sys=nothing)
    SciMLFunction{1, 2, typeof(f), typeof(mass_matrix), typeof(analytic), typeof(tgrad), typeof(jac), typeof(jvp), typeof(vjp), typeof(jac_prototype),
    typeof(sparsity), typeof(Wfact), typeof(Wfact_t), typeof(paramjac), typeof(syms), typeof(indepsym), typeof(paramsyms),
    typeof(observed), typeof(colorvec), typeof(sys)}(f, mass_matrix, analytic, tgrad, jac, jvp, vjp, jac_prototype,
        sparsity, Wfact, Wfact_t, paramjac, syms, indepsym, paramsyms,
        observed, colorvec, sys)
end

function foo(funcs::SciMLFunction, x)
    g = isnothing(funcs.jac) ? identity : funcs.jac
    h = isnothing(funcs.jvp) ? identity : funcs.jvp
    k = isnothing(funcs.vjp) ? identity : funcs.vjp
    funcs.f(g(h(k(x))))
end

struct MyFunction{F,N,T}
    f::F
    more_funcs::NamedTuple{N,T}
end

function MyFunction(f; kwargs...)
    # Extract the named tuple from the Base.Pairs to reduce
    # the type complexity.
    MyFunction(f, values(kwargs))
end

function foo(funcs::MyFunction, x)
    g = get(funcs.more_funcs, :jac, identity)
    h = get(funcs.more_funcs, :jvp, identity)
    k = get(funcs.more_funcs, :vjp, identity)
    funcs.f(g(h(k(x))))
end

Now let’s define an instance of each of the two types where none of the optional arguments are passed in:

sciml_funcs = SciMLFunction(x -> x + 1)
my_funcs = MyFunction(x -> x + 1)

Here is a screenshot of what the SciMLFunction stacktrace looks like:

Here is a screenshot of what the MyFunction stacktrace looks like:

If you compare frame 2 of the stacktraces, you can see that the printing of the type MyFunction is much shorter than the printing of the type SciMLFunction. I think that’s a considerable improvement.

The same technique could be applied to other types which contain a bunch of optional fields, like DEOptions.

4 Likes

Union splitting blows up compile times so we have to avoid that (and in fact, internally we turn off union splitting so this would be slow).

This doesn’t have error checking for typos. Incorrectly labelled functions like t_grad instead of tgrad are then just not found with has_property(x,:tgrad), so you get silent incorrect failures to use the function.

That was just a toy example. It’s pretty easy to address your concern.

struct MyFunction{F,N,T}
    f::F
    more_funcs::NamedTuple{N,T}
end

function MyFunction(
        f;
        mass_matrix=nothing, analytic=nothing, tgrad=nothing,
        jac=nothing, jvp=nothing, vjp=nothing,
        jac_prototype=nothing, sparsity=nothing, Wfact=nothing,
        Wfact_t=nothing, paramjac=nothing, syms=nothing,
        indepsym=nothing, paramsyms=nothing, observed=nothing, 
        colorvec=nothing, sys=nothing
    )
    # This particular syntax requires Julia 1.7.
    kwargs = (;
        mass_matrix, analytic, tgrad, jac, jvp, vjp,
        jac_prototype, sparsity, Wfact, Wfact_t, paramjac,
        syms, indepsym, paramsyms, observed, colorvec, sys
    )

    more_funcs = NamedTuple(
        k => v for (k, v) in pairs(kwargs) if v !== nothing
    )
    MyFunction(f, more_funcs)
end

Now we get the usual MethodError for incorrectly typed keyword argument names:

julia> MyFunction(x -> x + 1; t_grad = x -> x + 2)
ERROR: MethodError: no method matching MyFunction(::var"#57#59"; t_grad=var"#58#60"())
Closest candidates are:
  MyFunction(::Any; mass_matrix, analytic, tgrad, jac, jvp, vjp, jac_prototype, sparsity, Wfact, Wfact_t, paramjac, syms, indepsym, paramsyms, observed, colorvec, sys) at ~/projects/julia_misc/ode_function.jl:57 got unsupported keyword argument "t_grad"
  MyFunction(::F, ::NamedTuple{N, T}) where {F, N, T} at ~/projects/julia_misc/ode_function.jl:47 got unsupported keyword argument "t_grad"
Stacktrace:
 [1] kwerr(::NamedTuple{(:t_grad,), Tuple{var"#58#60"}}, ::Type, ::Function)
   @ Base ./error.jl:165
 [2] top-level scope
   @ REPL[24]:1
2 Likes

That is an interesting idea, and if a real solution wasn’t around the corner that would be okay. But it’s a very heavy way to simulate a parameterized type just because there’s no way to control stacktraces.

I take the reverse view. I would say a type with every field parameterized is a heavy way to represent a collection of optional configuration variables. A minimally parameterized struct makes sense for holding configuration variables if all the variables have sensible and usable default values. For a case like ODEFunction, where most of the configuration variables do not have usable default values, and most of the variables are left unspecified, it makes more sense to store the specified variables in a collection. The number of configuration variables that are set in the ODEFunction constructor can vary anywhere between 0 and 18. That to me sounds like a job for a collection, not a type with 18 parameters. (Actually 20 parameters since there are also the iip and specialize parameters.)

Actually, I think the most natural way to handle this sort of scenario is to just slurp and splat kwargs among the various functions involved. However, I’m not sure if that could be made to work without breaking the current API.

4 Likes

It could via getproperty overloading.

Note that this is type-unstable. Example:

julia> nt = (a=1, b=2, c=nothing);

julia> @btime NamedTuple(k=>v for (k,v) ∈ pairs($nt) if !isnothing(v))
  844.615 ns (15 allocations: 912 bytes)
(a = 1, b = 2)

Generally, using k=>v pairs to represent NamedTuples loses type information.

This, however, is type-stable:

julia> Base.filter(f, nt::NamedTuple{ks}) where ks = let ks=filter(f ∘ Base.Fix1(getindex, nt), ks)
           NamedTuple{ks}(map(Base.Fix1(getindex, nt), ks))
       end
       @btime filter(!isnothing, $nt)
  1.800 ns (0 allocations: 0 bytes)
(a = 1, b = 2)

Nice. Though if you’re only constructing a few ODEFunctions, 800 nanoseconds is still pretty fast.

By the way, the less obscure version of your function also has zero allocations:

function Base.filter(f, nt::NamedTuple{ks}) where ks
    new_ks = filter(k -> f(nt[k]), ks)
    NamedTuple{new_ks}(map(k -> nt[k], new_ks))
end
julia> nt = (a=1, b=2, c=nothing);

julia> @btime filter(!isnothing, $nt);
  2.552 ns (0 allocations: 0 bytes)

Don’t make your code obscure if the readable version runs just as fast. :wink:

(Of course we wouldn’t want to engage in type piracy on filter(f, ::NamedTuple), but that’s easily fixed.)

4 Likes

Would you like to open the PR?

Oh, I meant in this case we would just define a new function, like

trim(nt::NamedTuple) = nt[filter(k -> !isnothing(nt[k]), keys(nt))]
julia> nt = (a=1, b=2, c=nothing);

julia> trim(nt)
(a = 1, b = 2)

julia> @btime trim($nt);
  2.552 ns (0 allocations: 0 bytes)

But I suppose it probably would be nice to have filter on named tuples in Base. Feel free to open a PR if you like. :slight_smile:

2 Likes

@ChrisRackauckas try ConcreteStructs.jl with terse. It automatically avoids printing types in the stacktraces (it adds “quiet” types).

That has exactly the same issue.