Strip zygote frames from mutation error stack trace about zygote.jl HOT 1 OPEN

The story with Zygote stacktraces is more complex than described and could use a little explaining. If we use this stacktrace for illustration:

Stacktrace:
  [1] error(s::String)
    @ Base ./error.jl:35
  [2] _throw_mutation_error(f::Function, args::Vector{Any})
    @ Zygote ~/.julia/packages/Zygote/jxHJc/src/lib/array.jl:70

The first two stackframes are what you'd expect, common error reporting code. More interesting is the next two:

  [3] (::Zygote.var"#547#548"{Vector{Any}})(::Nothing)
    @ Zygote ~/.julia/packages/Zygote/jxHJc/src/lib/array.jl:89
  [4] (::Zygote.var"#2643#back#549"{Zygote.var"#547#548"{Vector{Any}}})(Δ::Nothing)
    @ Zygote ~/.julia/packages/ZygoteRules/M4xmc/src/adjoint.jl:72

As you've correctly identified, [3] has the actual rule we should look at. So what is [4]? That would be the rule machinery itself at https://github.com/FluxML/ZygoteRules.jl/blob/f9bf0e367fa259c5aa68f0e14ccbf2125d734bd6/src/adjoint.jl#L72. Not very helpful.

Now for the surprising revelation: there is actually no "interleaving of Zygote frames in this stacktrace". From [2] to [33], it's all Zygote. But how can that be when we have frames like this?

  [5] merge
    @ ./namedtuple.jl:371 [inlined]
  [6] (::Zygote.Pullback{Tuple{typeof(merge), @NamedTuple{}, Base.Generator{…}}, Any})(Δ::Nothing)
    @ Zygote ~/.julia/packages/Zygote/jxHJc/src/compiler/interface2.jl:0

Essentially, Zygote's generated functions can spoof line numbers from the original function, so that [5] merge frame is actually the same call as [6] (the AD-generated pullback).

I don't know why this was done. It was either intentional to help with looking up the original function since the genfunc code provides little info to work with, or the line info is sticking around by accident. This snippet also shows the limitations Zygote has around stackframe printing. Ideally, we'd want the call info of [5] with the file name and line number of [6] to make stacktraces would be shorter and cleaner. I'm assuming this is possible, but the main problem is that Zygote's internals are a PITA to work with (mostly because of IRTools, IMO).

from zygote.jl.