Repository: sbi-dev/sbi

Issue:
Add a "transform" method for the `NFLowsFlow` and `ZukoFlow` object
**Is your feature request related to a problem? Please describe.**
To create the normalizing flow version of the `LC2ST` diagnostic, I need access to the transform function of a flow that transforms data into the base distribution. These are different for `nflows` and `zuko`. A common method for both objects would make life easier.

**Describe the solution you'd like**
I would create a transform method for the base `DensityEstimator` class that takes data as input and transforms it into noise using the underlying normalizing flow.

**Describe alternatives you've considered**
...

**Additional context**
...


Discussion:
As `DensityEstimator` can be something different to a normalizing flow, I propose to create a method inside the `NFlowsFlow` and `ZukoFlow` objects directly

Resolving patch:
diff --git a/sbi/neural_nets/density_estimators/nflows_flow.py b/sbi/neural_nets/density_estimators/nflows_flow.py
--- a/sbi/neural_nets/density_estimators/nflows_flow.py
+++ b/sbi/neural_nets/density_estimators/nflows_flow.py
@@ -25,6 +25,55 @@ def embedding_net(self) -> nn.Module:
         r"""Return the embedding network."""
         return self.net._embedding_net
 
+    def inverse_transform(self, input: Tensor, condition: Tensor) -> Tensor:
+        r"""Return the inverse flow-transform of the inputs given a condition.
+
+        The inverse transform is the transformation that maps the inputs back to the
+        base distribution (noise) space.
+
+        Args:
+            input: Inputs to evaluate the inverse transform on of shape
+                    (*batch_shape1, input_size).
+            condition: Conditions of shape (*batch_shape2, *condition_shape).
+
+        Raises:
+            RuntimeError: If batch_shape1 and batch_shape2 are not broadcastable.
+
+        Returns:
+            noise: Transformed inputs.
+
+        Note:
+            This function should support PyTorch's automatic broadcasting. This means
+            the function should behave as follows for different input and condition
+            shapes:
+            - (input_size,) + (batch_size,*condition_shape) -> (batch_size,)
+            - (batch_size, input_size) + (*condition_shape) -> (batch_size,)
+            - (batch_size, input_size) + (batch_size, *condition_shape) -> (batch_size,)
+            - (batch_size1, input_size) + (batch_size2, *condition_shape)
+                                                  -> RuntimeError i.e. not broadcastable
+            - (batch_size1,1, input_size) + (batch_size2, *condition_shape)
+                                                  -> (batch_size1,batch_size2)
+            - (batch_size1, input_size) + (batch_size2,1, *condition_shape)
+                                                  -> (batch_size2,batch_size1)
+        """
+        self._check_condition_shape(condition)
+        condition_dims = len(self._condition_shape)
+
+        # PyTorch's automatic broadcasting
+        batch_shape_in = input.shape[:-1]
+        batch_shape_cond = condition.shape[:-condition_dims]
+        batch_shape = torch.broadcast_shapes(batch_shape_in, batch_shape_cond)
+        # Expand the input and condition to the same batch shape
+        input = input.expand(batch_shape + (input.shape[-1],))
+        condition = condition.expand(batch_shape + self._condition_shape)
+        # Flatten required by nflows, but now both have the same batch shape
+        input = input.reshape(-1, input.shape[-1])
+        condition = condition.reshape(-1, *self._condition_shape)
+
+        noise, _ = self.net._transorm(input, context=condition)
+        noise = noise.reshape(batch_shape)
+        return noise
+
     def log_prob(self, input: Tensor, condition: Tensor) -> Tensor:
         r"""Return the log probabilities of the inputs given a condition or multiple
         i.e. batched conditions.
diff --git a/sbi/neural_nets/density_estimators/zuko_flow.py b/sbi/neural_nets/density_estimators/zuko_flow.py
--- a/sbi/neural_nets/density_estimators/zuko_flow.py
+++ b/sbi/neural_nets/density_estimators/zuko_flow.py
@@ -34,6 +34,55 @@ def embedding_net(self) -> nn.Module:
         r"""Return the embedding network."""
         return self._embedding_net
 
+    def inverse_transform(self, input: Tensor, condition: Tensor) -> Tensor:
+        r"""Return the inverse flow-transform of the inputs given a condition.
+
+        The inverse transform is the transformation that maps the inputs back to the
+        base distribution (noise) space.
+
+        Args:
+            input: Inputs to evaluate the inverse transform on of shape
+                    (*batch_shape1, input_size).
+            condition: Conditions of shape (*batch_shape2, *condition_shape).
+
+        Raises:
+            RuntimeError: If batch_shape1 and batch_shape2 are not broadcastable.
+
+        Returns:
+            noise: Transformed inputs.
+
+        Note:
+            This function should support PyTorch's automatic broadcasting. This means
+            the function should behave as follows for different input and condition
+            shapes:
+            - (input_size,) + (batch_size,*condition_shape) -> (batch_size,)
+            - (batch_size, input_size) + (*condition_shape) -> (batch_size,)
+            - (batch_size, input_size) + (batch_size, *condition_shape) -> (batch_size,)
+            - (batch_size1, input_size) + (batch_size2, *condition_shape)
+                                                  -> RuntimeError i.e. not broadcastable
+            - (batch_size1,1, input_size) + (batch_size2, *condition_shape)
+                                                  -> (batch_size1,batch_size2)
+            - (batch_size1, input_size) + (batch_size2,1, *condition_shape)
+                                                  -> (batch_size2,batch_size1)
+        """
+
+        self._check_condition_shape(condition)
+        condition_dims = len(self._condition_shape)
+
+        # PyTorch's automatic broadcasting
+        batch_shape_in = input.shape[:-1]
+        batch_shape_cond = condition.shape[:-condition_dims]
+        batch_shape = torch.broadcast_shapes(batch_shape_in, batch_shape_cond)
+        # Expand the input and condition to the same batch shape
+        input = input.expand(batch_shape + (input.shape[-1],))
+        emb_cond = self._embedding_net(condition)
+        emb_cond = emb_cond.expand(batch_shape + (emb_cond.shape[-1],))
+
+        dists = self.net(emb_cond)
+        noise = dists.transform(input)
+
+        return noise
+
     def log_prob(self, input: Tensor, condition: Tensor) -> Tensor:
         r"""Return the log probabilities of the inputs given a condition or multiple
         i.e. batched conditions.


Test patch: