Rémi Louf
Rémi Louf
#305 is related
Hi @karm-patel. Thank you for you interest in `blackjax` and your willingness to contribute! I think this function would be better implemented in `arviz`; we do not want to add...
@karm-patel Feel free to open a PR for the `sample` loop, I'd be happy to take a look.
I'll close this issue for now. You can open a new one for sampling loops if that's something you're interested in!
@karm-patel @murphyk We are considering opening a separate repo to create helper functions to wrap other PPLs's model and export the sampling information to ArviZ. Are you interested?
> Sounds great! I hope I could, but these days I'm occupied in my school studies and not getting the time for open source :( No worries, school takes at...
Thanks for contributing! Could you rebase your branch on `main`? There's a commit there fixing the error we're seeing in CI.
> Aesara functions do many other things besides calling the inner function. For example, they also perform updates of shared variables. At the very least shouldn't `jit_fn` not return a...
Comes from this line: https://github.com/aesara-devs/aesara/blob/84e69fc890434564757747fafc7646e54db6e14f/aesara/link/utils.py#L773 The original code did not have the `squeeze_output` keyword. I think one quick fix is to pass `squeeze_output` to `jax_funcify` when applied to a `FunctionGraph`...
It seems that fiddling with the JAX linker would lead to modifications in the `Function` class (which may be more appropriately called `CFunction`?). We shouldn't "special case" this class for...