use in TransformVariables.jl (seeking advice)

[tpapp]

I am refactoring TransformVariables.jl. I have a conceptual question which I think this package solves (I think), so I want to ask for advice.

For those unfamiliar with the package, it encodes uni- and multivariate domain transformations, mostly for use in Bayesian estimation with samplers that work on an unconstrained Euclidean space. When calculating a log likelihood, this requires a correction by the log of the absolute value of the determinant of the Jacobian.

The current API is

y, lj = transform_and_logjac(transformation, x)

mapping x to y, with lj as the log abs det J.

This is cumbersome, as the user has to keep and accumulate ljs. Since those statements are usually followed by a logpdf(some_distribution, y) somewhere, I thought of returning a

struct ValueWithLogjac{T,L}
    y::T
    lj::L
end

wrapper, defining

Base.logpdf(distribution, v::ValueWithLogjac) = logpdf(distribution, v.y) + v.lj

but that would require loading Distributions.jl, a heavy dependency. Now JuliaStats/Distributions.jl#1139 would solve that but it seems not to be going anywhere. However, I could just do

DensityInterface.logdensityof(object, v::ValueWithLogjac) = DensityInterface.logdensityof(distribution, v.y) + v.lj

which works automatically for types in Distributions.jl, and is lightweight.

Does this conform with the intended usage of this package? Or should I try factoring out a DistributionsCore.jl again? Thoughts are appreciated.

[oschulz]

I think this could be very useful (though I think a DistributionsCore would be valuable in addition anyway, I often wished we had something like that), but there's a dispatch problem (see below for a proposed solution).

One question is, how would we best mesh it with ChangesOfVariables, which currently returns tuples. I'm not sure that should change, it would be quite a breaking change. It would allow though to support with_logabsdet_jacobian(f, x::ValueWithLogjac)::ValueWithLogjac. ChangesOfVariables already does handle function composition explicitly, but this would establish an underlying mechanism for it.

So a type ValueWithLADJ would do for log-abs-det-Jacobian propagation what ForwardDiff.Dual does for derivative-propagation, in a way. Is there an established mathematical framework for this, like there is for dual numbers?

But ... the major dispatch problem: we often do (and need to) specialize logdensityof(obj::SomeType, x::Any) and with_logabsdet_jacobian(f::SomeType, x::Any): For things like pushforward measures, function chains and many other things, these two need to be specialized for any x, the type of x isn't known. So specializing logdensityof(obj::Any, v::ValueWithLogjac) would land us in ambiguity hell.

So we'll definitely need to separate the function that we'll specialize based on the object/function type from the function we'll specialize based on the value type. One could add a logdensityof_impl for the first case, but that would require changing a lot of code using these two packages.

I think the best approach could be to introduce two new functions like

# In DensityInterface:

transformed_logdensity(obj::Any, x::Any) = logdensityof(obj, x)

transformed_logdensity(obj::Any, x::ValueWithLADJ) = logdensityof(obj, x.value) + x.ladj

and

# In ChangesOfVariables:

_to_ValueWithLADJ(x::NoLogAbsDetJacobian) = x
_to_ValueWithLADJ(x::Tuple{Any,Any}) = ValueWithLADJ(x[1], x[2])

propagate_ladj(f::Any, x::Any) = _to_ValueWithLADJ(with_logabdet_jacobian(f, x))

_to_ValueWithLADJ(x::NoLogAbsDetJacobian, prev_ladj) = x
_to_ValueWithLADJ(x::Tuple{Any,Any}, prev_ladj) = ValueWithLADJ(x[1], x[2] + prev_ladj)

propagate_ladj(f::Any, x::ValueWithLADJ) = _to_ValueWithLADJ(with_logabdet_jacobian(f, x.value), x.ladj)

Question is, in which package should we define ValueWithLADJ?

Type and function names would be up for discussion, of course.

@devmotion what are your thoughts on this?

[oschulz]

Gentle ping @devmotion .

[devmotion]

Oops sorry, missed the first ping.

So specializing logdensityof(obj::Any, v::ValueWithLogjac) would land us in ambiguity hell.

That was my main concern when reading the proposal.

Now JuliaStats/Distributions.jl#1139 would solve that

That would lead to the same ambiguity issues it seems.

---

Given these complications, I'm generally a bit unsure whether one should move on with the proposal in TransformVariables. I think currently I'm lacking a good motivation for such a change as the tuple return values have worked flawlessly for me in practice.

[oschulz]

@devmotion sure, but what would you think about a "higher-level" API like transformed_logdensity and propagate_ladj? We could add that without changing the Tuple return value of with_logabsdet_jacobian, which it would use under the hood.