TuringLang / Bijectors.jl

Implementation of normalising flows and constrained random variable transformations
https://turinglang.org/Bijectors.jl/
MIT License
196 stars 32 forks source link

Allow different types for upper and lower bound in `Logit` #305

Closed torfjelde closed 2 months ago

torfjelde commented 2 months ago

Currently the lower bound and upper bound are forced to be the same type, which removes the possibility of computing gradients wrt. just one of the parameter.

This PR just removes this restriction.

devmotion commented 2 months ago

This should not be required for AD - most distributions only have a single parameter type. I also think it won't be (significantly) more efficient - forward computations will use duals anyway if one parameter is a dual number, and backward computation will stop once they encounter a number that was promoted to a tracked one in the forward pass. Two types means more type parameters though and hence more stress on the compiler and also a reader/developer who tries to reason about return types.

So I'm not fully convinced by the change 😉

torfjelde commented 2 months ago

Good point :thinking: But doesn't converting both into, say, vector of duals lead to perf implications since subsequent computations then also computes derivatives wrt. this?

devmotion commented 2 months ago

Subsequent computations will be duals anyway if one of the fields are duals. So the only difference is that in one case the subsequent computation f(a, b) is started with both a and b being duals (but one of them having 0 partials) and in the other case one of them is a dual and the other one is non-dual but will be promoted to duals with 0 partials in every computation inside of f involving duals.

torfjelde commented 2 months ago

Only every computation in which they're both related, no? For example, in the computation (b - x) / (b - a), if only a is duals you'd unnecessarily also compute partials through b - x ?

But yeah, I do agree that it probably doesn't matter that much. Happy to accept a PR with proper promotion instead :+1: