Vilin97

"reset a leaf"
function reset!(leaf::Leaf{A, S}) where {A <: Optimisers.Adam, S}
    leaf.state[1] .= 0
    leaf.state[2] .= 0
    leaf.state = (leaf.state[1], leaf.state[2], leaf.rule.beta)
    nothing
end
"reset optimiser state"
function reset!(state::NamedTuple{(:layers,), L}) where {L}
    for layer in state.layers

Vilin97

using DiffEqJump, Random, BenchmarkTools, DataStructures
const MINJUMPRATE = 2.0^exponent(1e-12)

function pq_setup(rates)
    pqdata = randexp(length(rates))./rates
    MutableBinaryMinHeap(pqdata)
end

function pt_setup!(pt, rates)
    DiffEqJump.reset!(pt)

Vilin97

\DeclarePairedDelimiter{\norm}{||}{||_1}
\providecommand{\floor}[1]{\left \lfloor #1 \right \rfloor }
\newcommand{\shortexactseq}[5]{#1\to #2 \hookrightarrow #3 \twoheadrightarrow #4 \to #5}
\newcommand{\rad}[1]{\sqrt{#1}}
\newcommand{\set}[1]{\left\{ #1 \right\}}
\newcommand{\gen}[1]{\langle #1 \rangle}
\newcommand{\restr}[1]{|_{#1}}
\newcommand{\sepdeg}[2]{[{#1}:{#2}]_{\rm sep}}

\renewcommand{\Pr}{\mathbb{P}}