Data adapter: Keep track of the jacobian for variable transforms

[Kucharssim]

Given that the data adapter is now used for variable transforms (for e.g., implementing parameter constraints or scaling), it would be nice to keep track of the jacobian of the transform, so that if the user wants to use the log_prob method, it is possible to automatically obtain the log prob on the "constrained/data space" instead of the "unconstrained/network space".

For custom transforms, the user would have to supply corresponding log_det_jacobian function, but of course that should be only optional to those cases where the user wants to use log_prob in the first place.

I know we talked about this with @LarsKue and @paul-buerkner already at different time points, just making it an issue so that we can keep track of it here.

[stefanradev93]

Hi Simon, I like the idea! Would you like to try a draft implementation a few pull requests down the line?

[Kucharssim]

Sure, I can give it a go :)

[LarsKue]

@Kucharssim any update on this? We still think it would be a nice feature to have.

[Kucharssim]

Sorry for sitting on this.

I have made progress on this at some point and then let it get stale, I have now rebased it on the current dev branch. You can check the functionality here: https://github.com/Kucharssim/BayesFlow/tree/jacobian

I have not done the jacobian for the inverse (though it shouldn't be that difficult for the most part), and there are still some transforms for which I haven't implemented the jacobian (e.g., NumpyTransform). But the basic use case works on that branch.

Since I am not a pythonista, I am sure I didn't implement this in the cleanest way possible. If you have some tips how to simplify the code, that would be very much appreciated.

I can in the meantime look into implementing the missing features, and then I can make a PR for this. Does that sound ok?

[LarsKue]

Hi @Kucharssim, I checked out your fork and from my side it looks fine. In Normalizing Flows, the convention is to compute the log_det along with the forward call for efficiency, but I can see how that is not as relevant for Adapter Transforms since they are cheap, and also more difficult to implement with the current state of the Adapter.

That said, if you want to copy this pattern, feel free to! Please make a PR whenever you deem the feature ready for review.

[Kucharssim]

Yeah, I saw that pattern, but it seems to me that the majority of adapter calls happen during training, for which we typically do not need the log_det_jac at all, only the forward transform (whereas with normalizing flows, we need the log_det_jac for training). So computing the jacobians along with the forward transforms seems to me a bit wasteful. But I would be happy to refactor this if you have a strong preference for it.

[LarsKue]

No strong preference on my side.

[vpratz]

Closed by #419, if I read it correctly.