Repository: sbi-dev/sbi

Issue:
Compatibility of LazyTransform and AffineTransform | Unconditional
In e9d8dff7e65513ed994ec0117e70ea06cb2e2431 (PR https://github.com/sbi-dev/sbi/pull/1045) we currently ignore that `LazyTransform` and `AffineTransform | Unconditional` are incompatible typewise but this will probably be an issue which will come up more often. 
For example right now in #1068 (https://github.com/sbi-dev/sbi/actions/runs/8364627831/job/22900410899): when passing `transform: AffineTransform | Unconditional` into `Flow(transform: LazyTransform, ...)` 
 ```python
neural_net = zuko.flows.Flow(transform, flow_built.base)
```


Discussion:
I think we could fix this if we return a `LazyTransform` here:

https://github.com/sbi-dev/sbi/blob/89c2bde9f5f65b61e42941858a675bf6e01858e3/sbi/utils/sbiutils.py#L178-L186

At least if the backend is `zuko`. @gmoss13 Maybe we can have a look into this together?
I think this is an issue on the side of Zuko, as they use `Unconditional` effectively as a `LazyTransform`, but don't explicitly inherit from `LazyTransform`. I will write a wrapper we can use, and in the meantime let the Zuko devs know about this!

Resolving patch:
diff --git a/sbi/neural_nets/flow.py b/sbi/neural_nets/flow.py
--- a/sbi/neural_nets/flow.py
+++ b/sbi/neural_nets/flow.py
@@ -14,12 +14,12 @@
     rational_quadratic,  # pyright: ignore[reportAttributeAccessIssue]
 )
 from torch import Tensor, nn, relu, tanh, tensor, uint8
-from zuko.flows import LazyTransform
 
 from sbi.neural_nets.density_estimators import NFlowsFlow, ZukoFlow
 from sbi.utils.sbiutils import (
     standardizing_net,
     standardizing_transform,
+    standardizing_transform_zuko,
     z_score_parser,
 )
 from sbi.utils.torchutils import create_alternating_binary_mask
@@ -501,15 +501,12 @@ def build_zuko_maf(
             residual=residual,
         )
 
-    transforms: Union[Sequence[LazyTransform], LazyTransform]
-    transforms = maf.transform.transforms  # pyright: ignore[reportAssignmentType]
+    transforms = maf.transform
     z_score_x_bool, structured_x = z_score_parser(z_score_x)
     if z_score_x_bool:
-        # transforms = transforms
         transforms = (
-            *transforms,
-            # Ideally `standardizing_transform` would return a `LazyTransform` instead of ` AffineTransform | Unconditional`, maybe all three are compatible
-            standardizing_transform(batch_x, structured_x, backend="zuko"),  # pyright: ignore[reportAssignmentType]
+            transforms,
+            standardizing_transform_zuko(batch_x, structured_x),
         )
 
     z_score_y_bool, structured_y = z_score_parser(z_score_y)
diff --git a/sbi/utils/sbiutils.py b/sbi/utils/sbiutils.py
--- a/sbi/utils/sbiutils.py
+++ b/sbi/utils/sbiutils.py
@@ -15,11 +15,18 @@
 from pyro.distributions import Empirical
 from torch import Tensor, ones, optim, zeros
 from torch import nn as nn
-from torch.distributions import Distribution, Independent, biject_to, constraints
+from torch.distributions import (
+    AffineTransform,
+    Distribution,
+    Independent,
+    biject_to,
+    constraints,
+)
 
 from sbi import utils as utils
 from sbi.sbi_types import TorchTransform
 from sbi.utils.torchutils import atleast_2d
+from sbi.utils.zukoutils import UnconditionalLazyTransform
 
 
 def warn_if_zscoring_changes_data(x: Tensor, duplicate_tolerance: float = 0.1) -> None:
@@ -140,9 +147,8 @@ def standardizing_transform(
     batch_t: Tensor,
     structured_dims: bool = False,
     min_std: float = 1e-14,
-    backend: str = "nflows",
-) -> Union[transforms.AffineTransform, zuko.flows.Unconditional]:
-    """Builds standardizing transform
+) -> transforms.AffineTransform:
+    """Builds standardizing transform for nflows
 
     Args:
         batch_t: Batched tensor from which mean and std deviation (across
@@ -157,7 +163,59 @@ def standardizing_transform(
     Returns:
         Affine transform for z-scoring
     """
+    t_mean, t_std = z_standardization(batch_t, structured_dims, min_std)
+    return transforms.AffineTransform(shift=-t_mean / t_std, scale=1 / t_std)
 
+
+def standardizing_transform_zuko(
+    batch_t: Tensor,
+    structured_dims: bool = False,
+    min_std: float = 1e-14,
+) -> zuko.flows.LazyTransform:
+    """Builds standardizing transform for Zuko flows
+
+    Args:
+        batch_t: Batched tensor from which mean and std deviation (across
+            first dimension) are computed.
+        structured_dim: Whether data dimensions are structured (e.g., time-series,
+            images), which requires computing mean and std per sample first before
+            aggregating over samples for a single standardization mean and std for the
+            batch, or independent (default), which z-scores dimensions independently.
+        min_std:  Minimum value of the standard deviation to use when z-scoring to
+            avoid division by zero.
+
+    Returns:
+        Affine transform for z-scoring
+    """
+    t_mean, t_std = z_standardization(batch_t, structured_dims, min_std)
+    return UnconditionalLazyTransform(
+        AffineTransform,
+        loc=-t_mean / t_std,
+        scale=1 / t_std,
+        buffer=True,
+    )
+
+
+def z_standardization(
+    batch_t: Tensor,
+    structured_dims: bool = False,
+    min_std: float = 1e-14,
+) -> [Tensor, Tensor]:
+    """Computes mean and standard deviation for z-scoring
+
+    Args:
+        batch_t: Batched tensor from which mean and std deviation (across
+            first dimension) are computed.
+        structured_dim: Whether data dimensions are structured (e.g., time-series,
+            images), which requires computing mean and std per sample first before
+            aggregating over samples for a single standardization mean and std for the
+            batch, or independent (default), which z-scores dimensions independently.
+        min_std:  Minimum value of the standard deviation to use when z-scoring to
+            avoid division by zero.
+
+    Returns:
+        Mean and standard deviation for z-scoring
+    """
     is_valid_t, *_ = handle_invalid_x(batch_t, True)
 
     if structured_dims:
@@ -175,18 +233,8 @@ def standardizing_transform(
         t_std = torch.std(batch_t[is_valid_t], dim=0)
         t_std[t_std < min_std] = min_std
 
-    if backend == "nflows":
-        return transforms.AffineTransform(shift=-t_mean / t_std, scale=1 / t_std)
-    elif backend == "zuko":
-        return zuko.flows.Unconditional(
-            zuko.transforms.MonotonicAffineTransform,
-            shift=-t_mean / t_std,
-            scale=1 / t_std,
-            buffer=True,
-        )
-
-    else:
-        raise ValueError("Invalid backend. Use 'nflows' or 'zuko'.")
+    # Return mean and std for z-scoring.
+    return t_mean, t_std
 
 
 class Standardize(nn.Module):
diff --git a/sbi/utils/zukoutils.py b/sbi/utils/zukoutils.py
new file mode 100644
--- /dev/null
+++ b/sbi/utils/zukoutils.py
@@ -0,0 +1,7 @@
+from zuko.flows import LazyTransform, Unconditional
+
+
+# This is a temporary wrapper for the Unconditional class in zuko.
+# Avoids pyright error of zuko Flows requiring a LazyTransform as input.
+class UnconditionalLazyTransform(Unconditional, LazyTransform):
+    pass


Test patch: