Replay the MDN's i.i.d. density as TorchScript instead of driving it from R
Source:R/likelihood.R
mdn_trace_cache.RdEvery operation in mdn_iid_blocks() crosses from R into libtorch, and at
MCMC batch sizes that crossing costs far more than the arithmetic behind it:
about 0.2 ms each, thirty of them per evaluation, against a few hundred
microseconds of actual work. torch::jit_trace() records the same
computation once and replays it inside libtorch, so an evaluation costs one
crossing rather than thirty.
Arguments
- de, xt, max_batch
As in
mdn_iid_blocks(), with the observations already a tensor.- eager
The evaluator to check each trace against, and to fall back to.
- warmup
Calls to serve eagerly before recording anything. No caller overrides the default today, but it is a real tuning knob – how many evaluations tracing costs before it pays for itself – that a future caller would plausibly want to change, so it stays a parameter.
Value
function(theta) returning a traced function for that many rows, or
NULL when the eager path should be used. NULL if tracing is switched
off.
Details
It is the same code either way. Tracing runs the eager path and records what
it did, so there is no separate implementation to keep in step with the
eager one – that claim holds only for the traced path relative to
mdn_log_prob_tensor(). The MDN density still has three implementations in
total, one per runtime: mdn_log_prob_tensor() (the eager training path),
mdn_mixture()/mdn_chunk_lp() (the i.i.d. fast path used here), and
stan_fn_mdn() (the generated Stan code, R/stan.R). That is by design,
not drift: three runtimes need three implementations, and the tests pin
them to each other numerically.
Three things make this a shortcut rather than the path. Recording a trace
costs several evaluations' worth of time, so nothing is recorded until the
evaluator has been called warmup times and it is clear this is a loop
rather than a one-off; a single log_lik() call should not pay for a
compiler. A trace fixes the shapes it was recorded at, so there is one per
number of parameter rows, checked against the eager result before anything
uses it. And it is only worth recording when the whole observation set fits
in one chunk, since otherwise the graph unrolls the chunk loop. Failing any
of these is not an error: the caller falls back to the eager path, which is
why NULL is a perfectly good answer here.
Set options(neuralsbi.jit = FALSE) to skip tracing entirely.