Getting rid of need for get_log_abs_det_jacobian function

[janfb]

Long story short, this PR proposes a fix for the problem below by constructing all transforms such that the sum over the second dimension of the parameters Tensor passed to the ladj method.

At the moment we have a helper function get_log_abs_det_jacobian (ladj) to make sure the ladj is returned summed over the parameter dimension. I looked into this a bit and understood why sometimes the sum is performed and sometimes not.

The "raw" transforms from torch.distributions.transforms assume event_dim=0 when initialised, so the transforms constructed here:

sbibm/sbibm/utils/pyro.py

Lines 99 to 102 in 89755d2

	if automatic_transform_enabled:
	transforms[name] = biject_to(fn.support).inv
	else:
	transforms[name] = dist.transforms.identity_transform

will all have event_dim=0.
Thus, given a parameter with shape (batch_size, dim_parameters), accumulating the ladj to get a transform-corrected log_prob here:

sbibm/sbibm/utils/pyro.py

Lines 211 to 214 in 89755d2

	for name, t in self.transforms.items():
	log_joint -= get_log_abs_det_jacobian(
	t, params_constrained[name], params[name]
	)

would result in a ladj shape of (batch_size, dim_parameters) instead of (batch_size, ) (unless we use the get_log_abs_det_jacobian as we do it currently).

But I think there is a cleaner solution. We noticed before that in some cases the sum over the second dimension is indeed calculated. This is the case when the transform is defined as an IndependentTransform with reinterpreted_batch_ndims=1, and this happens automatically when calling biject_to on a BoxUniform prior because this prior has reinterpreted batch dims already.

With this fix we would not need get_log_abs_det_jacobian anymore. However, we first have to check it is bullet proof, e.g., what happens if the prior has more than two dimensions, is this case even allowed in sbibm?

[janfb]

close in favour of #27