inverse for Transforms

Author: Kucharssim

bayesflow/adapters/transforms/concatenate.py

@@ -128,11 +128,26 @@ def extra_repr(self) -> str:
         return result
 
     def log_det_jac(
-        self, data: dict[str, np.ndarray], log_det_jac: dict[str, np.ndarray], *, strict: bool = False, **kwargs
+        self,
+        data: dict[str, np.ndarray],
+        log_det_jac: dict[str, np.ndarray],
+        *,
+        strict: bool = False,
+        inverse: bool = False,
+        **kwargs,
     ) -> dict[str, np.ndarray]:
         # copy to avoid side effects
         log_det_jac = log_det_jac.copy()
 
+        if inverse:
+            if log_det_jac.get(self.into) is not None:
+                raise ValueError(
+                    "Cannot obtain an inverse jacobian of concatenation. "
+                    "Transform your variables before you concatenate."
+                )
+
+            return log_det_jac
+
         required_keys = set(self.keys)
         available_keys = set(log_det_jac.keys())
         common_keys = available_keys & required_keys

bayesflow/adapters/transforms/drop.py

@@ -55,5 +55,5 @@ def inverse(self, data: dict[str, any], **kwargs) -> dict[str, any]:
     def extra_repr(self) -> str:
         return "[" + ", ".join(map(repr, self.keys)) + "]"
 
-    def log_det_jac(self, data: dict[str, any], log_det_jac: dict[str, any], **kwargs):
-        return self.forward(data=log_det_jac)
+    def log_det_jac(self, data: dict[str, any], log_det_jac: dict[str, any], inverse: bool = False, **kwargs):
+        return self.inverse(data=log_det_jac) if inverse else self.forward(data=log_det_jac)

bayesflow/adapters/transforms/keep.py

@@ -66,5 +66,5 @@ def inverse(self, data: dict[str, any], **kwargs) -> dict[str, any]:
     def extra_repr(self) -> str:
         return "[" + ", ".join(map(repr, self.keys)) + "]"
 
-    def log_det_jac(self, data: dict[str, any], log_det_jac: dict[str, any], **kwargs):
-        return self.forward(data=log_det_jac)
+    def log_det_jac(self, data: dict[str, any], log_det_jac: dict[str, any], inverse: bool = False, **kwargs):
+        return self.inverse(data=log_det_jac) if inverse else self.forward(data=log_det_jac)

bayesflow/adapters/transforms/rename.py

@@ -68,5 +68,5 @@ def inverse(self, data: dict[str, any], *, strict: bool = False, **kwargs) -> di
     def extra_repr(self) -> str:
         return f"{self.from_key!r} -> {self.to_key!r}"
 
-    def log_det_jac(self, data: dict[str, any], log_det_jac: dict[str, any], **kwargs):
-        return self.forward(data=log_det_jac, strict=False)
+    def log_det_jac(self, data: dict[str, any], log_det_jac: dict[str, any], inverse: bool = False, **kwargs):
+        return self.inverse(data=log_det_jac) if inverse else self.forward(data=log_det_jac, strict=False)

bayesflow/adapters/transforms/transform.py

@@ -36,6 +36,6 @@ def extra_repr(self) -> str:
         return ""
 
     def log_det_jac(
-        self, data: dict[str, np.ndarray], log_det_jac: dict[str, np.ndarray], **kwargs
+        self, data: dict[str, np.ndarray], log_det_jac: dict[str, np.ndarray], inverse: bool = False, **kwargs
     ) -> dict[str, np.ndarray]:
         return log_det_jac