inverse for elementwise

Author: Kucharssim

bayesflow/adapters/transforms/constrain.py

@@ -207,6 +207,8 @@ def inverse(self, data: np.ndarray, **kwargs) -> np.ndarray:
         # inverse means network space -> data space, so constrain the data
         return self.constrain(data)
 
-    def log_det_jac(self, data: np.ndarray, **kwargs) -> np.ndarray:
+    def log_det_jac(self, data: np.ndarray, inverse: bool = False, **kwargs) -> np.ndarray:
         ldj = self.ldj(data)
+        if inverse:
+            ldj = -ldj
         return np.sum(ldj, axis=tuple(range(1, ldj.ndim)))

bayesflow/adapters/transforms/elementwise_transform.py

@@ -25,5 +25,5 @@ def forward(self, data: np.ndarray, **kwargs) -> np.ndarray:
     def inverse(self, data: np.ndarray, **kwargs) -> np.ndarray:
         raise NotImplementedError
 
-    def log_det_jac(self, data: np.ndarray, **kwargs) -> np.ndarray | None:
+    def log_det_jac(self, data: np.ndarray, inverse: bool = False, **kwargs) -> np.ndarray | None:
         return None

bayesflow/adapters/transforms/log.py

@@ -49,6 +49,8 @@ def get_config(self) -> dict:
             "p1": serialize(self.p1),
         }
 
-    def log_det_jac(self, data: np.ndarray, **kwargs) -> np.ndarray:
+    def log_det_jac(self, data: np.ndarray, inverse: bool = False, **kwargs) -> np.ndarray:
         ldj = -np.log(data)
+        if inverse:
+            ldj = -ldj
         return np.sum(ldj, axis=tuple(range(1, ldj.ndim)))

bayesflow/adapters/transforms/scale.py

@@ -26,7 +26,9 @@ def forward(self, data: np.ndarray, **kwargs) -> np.ndarray:
     def inverse(self, data: np.ndarray, **kwargs) -> np.ndarray:
         return data / self.scale
 
-    def log_det_jac(self, data: np.ndarray, **kwargs) -> np.ndarray:
+    def log_det_jac(self, data: np.ndarray, inverse: bool = False, **kwargs) -> np.ndarray:
         ldj = np.log(np.abs(self.scale))
         ldj = np.full(data.shape, ldj)
+        if inverse:
+            ldj = -ldj
         return np.sum(ldj, axis=tuple(range(1, ldj.ndim)))

bayesflow/adapters/transforms/sqrt.py

@@ -26,6 +26,8 @@ def from_config(cls, config: dict, custom_objects=None) -> "Sqrt":
     def get_config(self) -> dict:
         return {}
 
-    def log_det_jac(self, data: np.ndarray, **kwargs) -> np.ndarray:
+    def log_det_jac(self, data: np.ndarray, inverse: bool = False, **kwargs) -> np.ndarray:
         ldj = -0.5 * np.log(data) + 0.5
+        if inverse:
+            ldj = -ldj
         return np.sum(ldj, axis=tuple(range(1, ldj.ndim)))

bayesflow/adapters/transforms/standardize.py

@@ -130,9 +130,11 @@ def inverse(self, data: np.ndarray, **kwargs) -> np.ndarray:
 
         return data * std + mean
 
-    def log_det_jac(self, data, **kwargs) -> np.ndarray:
+    def log_det_jac(self, data, inverse: bool = False, **kwargs) -> np.ndarray:
         if self.std is None:
             return None
         std = np.broadcast_to(self.std, data.shape)
         ldj = np.log(np.abs(std))
+        if inverse:
+            ldj = -ldj
         return np.sum(ldj, axis=tuple(range(1, ldj.ndim)))