flowjax
flowjax copied to clipboard
danielward27
I am looking into a modification of a regular masked autoregressive flow where the base distribution is an N-dimensional uniform and the first variable does not get transformed, while the...
Sphinx doctest fails for a couple reasons. When running doctest 1. Sphinx seems to stringify types which is incompatible with equinox (cannot use stringified abstract annotations). 2. We set ``typing.GENERATING_DOCUMENTATION...
If you're building on `equinox`, might as well use `jaxtyping` for more detailed type hints, including shape information.
Instead of permutations, one can consider learnable linear transformations (See section 3.2 of the Papamakarios review.). As far as I understand it, in theory this could allow us to learn...
I am interested in using Optax's [stochastic gradient estimators](https://optax.readthedocs.io/en/latest/api/stochastic_gradient_estimators.html) and [control variates](https://optax.readthedocs.io/en/latest/api/control_variates.html) with FlowJAX. In particular, I am interested in compatibility with the reparameterization gradient (aka pathwise estimator). The reparameterization...
Clarify why constrained flows are needed. In my opinion, this example does not sufficiently motivate why constrained flows are needed. This makes it confusing for a new user, especially since...
Very nice package :-) While playing around with different optimization objectives I ran into an autodiff issue. The following always returns exactly zeros, which I think isn't correct. This might...
