Bugfix ranks for various shapes 3

Author: stefanradev93

bayesflow/coupling_networks.py

@@ -436,9 +436,10 @@ def _semantic_spline_parameters(self, parameters):
         """
 
         shape = tf.shape(parameters)
-        if len(shape) == 2:
+        rank = len(shape)
+        if rank == 2:
             new_shape = (shape[0], self.dim_out, -1)
-        elif len(shape) == 3:
+        elif rank == 3:
             new_shape = (shape[0], shape[1], self.dim_out, -1)
         else:
             raise NotImplementedError("Spline flows can currently only operate on 2D and 3D inputs!")

bayesflow/helper_networks.py

@@ -112,7 +112,7 @@ def call(self, target, condition, **kwargs):
 
         # Handle 3D case for a set-flow and repeat condition over
         # the second `time` or `n_observations` axis of `target``
-        if tf.rank(target) == 3 and tf.rank(condition) == 2:
+        if len(tf.shape(target)) == 3 and len(tf.shape(condition)) == 2:
             shape = tf.shape(target)
             condition = tf.expand_dims(condition, 1)
             condition = tf.tile(condition, [1, shape[1], 1])

bayesflow/inference_networks.py

@@ -188,7 +188,7 @@ def forward(self, targets, condition, **kwargs):
             condition_shape = tf.shape(condition)
 
             # Needs to be concatinable with condition
-            if tf.rank(condition) == 2:
+            if len(condition_shape) == 2:
                 shape_scale = (condition_shape[0], 1)
             else:
                 shape_scale = (condition_shape[0], condition_shape[1], 1)
@@ -201,7 +201,7 @@ def forward(self, targets, condition, **kwargs):
                 noise_scale = tf.zeros(shape=shape_scale) + self.soft_low
 
             # Perturb data with noise (will broadcast to all dimensions)
-            if len(shape_scale) == 2 and tf.rank(targets) == 3:
+            if len(shape_scale) == 2 and len(target_shape) == 3:
                 targets += tf.expand_dims(noise_scale, axis=1) * tf.random.normal(shape=target_shape)
             else:
                 targets += noise_scale * tf.random.normal(shape=target_shape)
@@ -228,7 +228,7 @@ def inverse(self, z, condition, **kwargs):
         if self.soft_flow and condition is not None:
             # Needs to be concatinable with condition
             shape_scale = (
-                (condition.shape[0], 1) if tf.rank(condition) == 2 else (condition.shape[0], condition.shape[1], 1)
+                (condition.shape[0], 1) if len(condition.shape) == 2 else (condition.shape[0], condition.shape[1], 1)
             )
             noise_scale = tf.zeros(shape=shape_scale) + 2.0 * self.soft_low