Open ColtAllen opened 1 month ago
Scan may be plenty fast in other backends: numba and jax, the first will be the default sometime in the future, and it's what it's used with nutpie. Jax is used for numpyro and blackjax. I would benchmark on those backends that before bothering to get rid of it.
Also for varied datasets (t_x very different across subjects) the non scan will probably be slower as it does a lot of useless computations. In the dense/ non scan way it will evaluate the worst case scenario (the biggest gap between T and t_x) for everyone even if it's only needed for 1 row out of 10000
ok! thanks for the input! I took the PR because I always wanna play with scan
, but we can close it and have other benchmarks. We can always come back and change it, as we have the code in a branch already.
logp
in theBetaGeoBetaBinom
distribution block contains an iterable currently serviced by aScan
frompytensor
. It's possible to refactor this so thatScan
is no longer needed:I compared both approaches in a dev notebook, and sans Scan is about 3x faster:
However, the above code requires modification because tests are failing with the returned
logp
values.