Extracting flatten/unflatten into lightweight package?

[oschulz]

In the spirit of creating lightweight interface-defining packages (see TuringLang/Bijectors.jl#199 which resulted in InverseFunctions.jl and ChangesOfVariables.jl):

While adding an interface-test utility to ChangesOfVariables.jl today, I needed a helper function _to_realvec_and_back - which does exactly what ParameterHandling.flatten does. And I needed it for exactly what #27 is requesting. :-)

It would be nice not to reinvent this, but ParameterHandling would of course be way to heavy a dependency for ChangesOfVariables - and even if it wasn't, we'd end up with a circular dependency since Bijectors will use ChangesOfVariables soon, so ParameterHandling will too.

I think the flatten/unflatten functionality of ParameterHandling does something very fundamental, and is orthogonal to it's variable transformation capabilities. Would you consider splitting it out into a lightweight package (basically the contents of flatten.jl)?

A truly lightweight recursive flatten/unflatten interface package could IMHO find use in many places in the ecosystem. flatten may be a bit too generic a name for the function (we probably have several flattens in the ecosystem), but how about flatten_and_back or so?

CC @willtebbutt, @paschermayr, @devmotion

Update: ChangesOfVariables.test_with_logabsdet_jacobian now has an additional optional argument to pass a transformation, which solves the dependency problem. Users will be able to use ParameterHandling for variable transformations during the test without a direct dependency between the two packages (after #27 is solved).

[willtebbutt]

How would you feel about depending on ParameterHandling once #42 goes in?

[oschulz]

How would you feel about depending on ParameterHandling once #42 goes in?

I think even with #42 it think it would still be too heavy - the idea is to keep ChangesOfVariables super-lightweight, so we can convince packages to depend on it, and the stuff in parameters.jl is really a different thing - here, I think, ParameterHandling would depend on ChangesOfVariables, not the other way round. And ParameterHandling is a more opinionated package than ChangesOfVariables.

But now that I think about it, since it looks like we need flatten/unflatten functionality in ChangesOfVariables, how would you like moving that there? It's not much more than I have in the draft of JuliaMath/ChangesOfVariables.jl#2 anyway, and I think it would fit into the theme of the package - it is a change of variables, after all (with a logabsdet-jacobian of 0, of course).

@willtebbutt and @devmotion, what do you think?

Update: ParameterHandling will likely depend on ChangesOfVariables indirectly in the future due to #42, since we hope to get support for ChangesOfVariables into LogExpFunctions (esp for logistic, logit etc.). So ChangesOfVariables must not depend on ParameterHandling.

[devmotion]

I'm not sure, it's not clear to me that this functionality should be included in JuliaMath/ChangesOfVariables.jl#2. I would prefer if such wrapping/reshaping/flattening etc. is handled by the user.

[oschulz]

Yes, the more I think about it, it's really a bit of a different thing. And we could use test_with_logabsdet_jacobian(f, x, torealvec_and_back, getjacobian), that would keep the whole thing out of the dependencies of ChangesOfVariables.

[oschulz]

I still think it's a very fundamental functionality that may deserve a central package independent from variable-domain transformations.

I wonder if we should add an interface InverseFunctions.with_inverse(f, x) with a default implementation of InverseFunctions.with_inverse(f, x) = (f(x), inverse(f)), to support use cases like flatten/unflatten that can't do inverse(f) without a value (they would specialize with_inverse)?

That would form a nice basis for a package that depends on InverseFunctions and defines a function realvec and with_inverse(::typeof(realvec), x). There are definitely use cases where one would only want to forward function and instantiating the back function would be wasteful (mem allocs). This way, we'd have an API than can do flatten and flatten-and-back.

[willtebbutt]

In principle I could get on board with this kind of thing.

My main concern is that if we choose to make this functionality more widely available it'll be difficult to get consensus on some semantics. For example, flatten(1) returns an empty vector and a closure. The rationale being that this is what you want if you're doing AD and 1 is a primal (as opposed to a tangent vector), because integers aren't differentiable. I can well imagine that someone somewhere would (quite reasonably) find this choice objectionable, and want different behaviour.

So I would want to generalise what we have here a little bit to allow different semantics depending on your situation. Maybe utilising a trait-based mechanism like the one used for RuleConfigs in ChainRulesCore?

So you would write something like

struct Flattener{T} end

struct KeepIntegers end

# These have different behaviour.
with_inverse(Flattener{}(), 1)
with_inverse(Flattener{KeepIntegers}(), 1) = [1.0]

in case you want to work around this stuff?

[oschulz]

Yes, thought some more about this today and also started to worry about integers. With the autodiff use case, as an extension to excluding integers, people may also want to exclude other parameters from the gradient calculation (if they are not relevant, to increase performance), so one might want to have a way to mark parameters as active/constant. That could quickly go beyond "non-opinionated central API" indeed ...

[willtebbutt]

so one might want to have a way to mark parameters as active/constant.

We actually already have a way to mark parameters as being constant (see fixed), although doing things the other way around (marking things as "active") is something that we don't currently have support for -- although we could almost certainly achieve using a trait-based system.

[oschulz]

I think it would be nice to explore these options in ParameterHandling (especially since it's going to me more lightweight soon). If we add InverseFunctions.with_inverse (JuliaMath/InverseFunctions.jl#6 - @devmotion?), your Flattener concept could be both powerful and user-friendly, I think.

A central "flatten-only" package could come later then, maybe? As soon as #27 is sorted out, people will already be able to ParameterHandling in ChangesOfVariables.test_with_logabsdet_jacobian (changed now), without ChangesOfVariables depending on ParameterHandling. And if we go for the with_inverse approach, I would simplify the rv_and_back = x -> (x, identity) option to transform = identity - much more elegant.

[oschulz]

We actually already have a way to mark parameters as being constant (see fixed)

Neat!

One advantage with the InverseFunctions.with_inverse approach would be that Flattener() itself would just return the value, not a tuple of value and back-function. So it could support ChangesOfVariables.with_logabsdet_jacobian. While Flattener() would have a trivial volume-element, it could be interesting for ValueFlattener() as a successor of value_flatten.

[oschulz]

@devmotion pointed me to https://github.com/JuliaDiff/FiniteDifferences.jl/blob/main/src/to_vec.jl

There's probably equivalents of this in several places in the ecosystem, internally. If we had a nice, lightweight central API for a to-real-vec-and-back, people could make their types support it, that would make a lot of AD a bit easier. It's handling integers and constant parameters though that makes it less straightforward than I thought initialy - but I think it could (and maybe should) still be decoupled from actual value transformations.

CC @mzgubic @oxinabox

[mzgubic]

Just to add that in the ChainRules ecosystem there is a vague plan for moving away from to_vec.

[oxinabox]

I think there is another copy of this general idea somewhere in ArrayInterface.jl?
cc @ChrisRackauckas