Bugfix of ranks 2

Author: stefanradev93

bayesflow/coupling_networks.py

@@ -436,9 +436,9 @@ def _semantic_spline_parameters(self, parameters):
         """
 
         shape = tf.shape(parameters)
-        if tf.rank(parameters) == 2:
+        if len(shape) == 2:
             new_shape = (shape[0], self.dim_out, -1)
-        elif tf.rank(parameters) == 3:
+        elif len(shape) == 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

@@ -228,7 +228,7 @@ def _forward(self, target):
         """Performs a learnable generalized permutation over the last axis."""
 
         shape = tf.shape(target)
-        rank = tf.rank(target)
+        rank = len(shape)
         log_det = tf.math.log(tf.math.abs(tf.linalg.det(self.W)))
         if rank == 2:
             z = tf.linalg.matmul(target, self.W)
@@ -241,7 +241,7 @@ def _inverse(self, z):
         """Un-does the learnable permutation over the last axis."""
 
         W_inv = tf.linalg.inv(self.W)
-        rank = tf.rank(z)
+        rank = len(tf.shape(z))
         if rank == 2:
             return tf.linalg.matmul(z, W_inv)
         return tf.tensordot(z, W_inv, [[rank - 1], [0]])
@@ -527,7 +527,7 @@ def call(self, x, **kwargs):
         # Example: Output dim is (batch_size, inv_dim) - > (batch_size, N, inv_dim)
         out_inv = self.invariant_module(x, **kwargs)
         out_inv = tf.expand_dims(out_inv, -2)
-        tiler = [1] * tf.rank(x)
+        tiler = [1] * len(shape)
         tiler[-2] = shape[-2]
         out_inv_rep = tf.tile(out_inv, tiler)