Follow-up to #27

sbibm/tasks/two_moons/task.py

@@ -165,7 +165,7 @@ def _get_transforms(
         *args,
         **kwargs: Any,
     ) -> Dict[str, Any]:
-        return {"parameters": torch.distributions.transforms.identity_transform}
+        return {"parameters": torch.distributions.transforms.IndependentTransform(torch.distributions.transforms.identity_transform, 1) }
 
     def _get_log_prob_fn(
         self,

sbibm/utils/pyro.py

@@ -98,13 +98,7 @@ def get_log_prob_fn(
         if automatic_transform_enabled:
             transforms[name] = biject_to(fn.support).inv
         else:
-            transforms[name] = dist.transforms.identity_transform
-        # Reinterpret batch dimensions of transform to get log abs det jac summed over
-        # parameter dimensions.
-        if not isinstance(transforms[name], IndependentTransform):
-            transforms[name] = IndependentTransform(
-                transforms[name], reinterpreted_batch_ndims=1
-            )
+            transforms[name] = dist.transforms.IndependentTransform(dist.transforms.identity_transform, 1)
 
     if implementation == "pyro":
         trace_prob_evaluator = TraceEinsumEvaluator(

tests/tasks/test_task_interface.py

@@ -101,3 +101,17 @@ def test_reference_posterior_not_called(task_name):
         reference_samples = task.get_reference_posterior_samples(num_observation=1)
 
     assert task is not None
+
+
+@pytest.mark.parametrize("task_name", [tn for tn in (all_tasks - julia_tasks)])
+def test_transforms_shapes(task_name, batch_size=5):
+    task = get_task(task_name)
+    prior = task.get_prior()
+    samples = prior(num_samples=batch_size)
+
+    transforms = task._get_transforms(True)["parameters"]
+
+    ladj_shape = transforms.log_abs_det_jacobian(transforms(samples), samples).shape
+    assert ladj_shape == torch.Size([batch_size])
+
+    assert transforms is not None