Use elementwise shift instead of elementwise affine in default LRS ar…

…chitecture
davidnabergoj · Nov 14, 2023 · 02f88ba · 02f88ba
1 parent f4805e9
commit 02f88ba
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/normalizing_flows/bijections/finite/autoregressive/architectures.py b/normalizing_flows/bijections/finite/autoregressive/architectures.py
@@ -141,13 +141,13 @@ class MaskedAutoregressiveLRS(BijectiveComposition):
     def __init__(self, event_shape, n_layers: int = 2, **kwargs):
         if isinstance(event_shape, int):
             event_shape = (event_shape,)
-        bijections = [ElementwiseAffine(event_shape=event_shape)]
+        bijections = [ElementwiseShift(event_shape=event_shape)]
         for _ in range(n_layers):
             bijections.extend([
                 ReversePermutation(event_shape=event_shape),
                 LRSForwardMaskedAutoregressive(event_shape=event_shape)
             ])
-        bijections.append(ElementwiseAffine(event_shape=event_shape))
+        bijections.append(ElementwiseShift(event_shape=event_shape))
         super().__init__(event_shape, bijections, **kwargs)