-
Notifications
You must be signed in to change notification settings - Fork 192
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Refactor logp
in BG/BB to remove Scan
#703
Comments
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 |
@ricardoV94 do you have a time estimate on when |
You can select other backends manually, don't need to wait for the default to change |
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.The text was updated successfully, but these errors were encountered: