Fix ndim -> rank bug

Author: stefanradev93

bayesflow/coupling_networks.py

@@ -436,9 +436,9 @@ def _semantic_spline_parameters(self, parameters):
         """
 
         shape = tf.shape(parameters)
-        if parameters.ndim == 2:
+        if tf.rank(parameters) == 2:
             new_shape = (shape[0], self.dim_out, -1)
-        elif parameters.ndim == 3:
+        elif tf.rank(parameters) == 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 target.ndim == 3 and condition.ndim == 2:
+        if tf.rank(target) == 3 and tf.rank(condition) == 2:
             shape = tf.shape(target)
             condition = tf.expand_dims(condition, 1)
             condition = tf.tile(condition, [1, shape[1], 1])
@@ -228,7 +228,7 @@ def _forward(self, target):
         """Performs a learnable generalized permutation over the last axis."""
 
         shape = tf.shape(target)
-        rank = target.ndim
+        rank = tf.rank(target)
         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 = z.ndim
+        rank = tf.rank(z)
         if rank == 2:
             return tf.linalg.matmul(z, W_inv)
         return tf.tensordot(z, W_inv, [[rank - 1], [0]])
@@ -402,11 +402,11 @@ def _initalize_parameters_data_dependent(self, init_data):
         """
 
         # 2D Tensor case, assume first batch dimension
-        if init_data.ndim == 2:
+        if tf.rank(init_data) == 2:
             mean = tf.math.reduce_mean(init_data, axis=0)
             std = tf.math.reduce_std(init_data, axis=0)
         # 3D Tensor case, assume first batch dimension, second number of observations dimension
-        elif init_data.ndim == 3:
+        elif tf.rank(init_data) == 3:
             mean = tf.math.reduce_mean(init_data, axis=(0, 1))
             std = tf.math.reduce_std(init_data, axis=(0, 1))
         # Raise other cases
@@ -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] * x.ndim
+        tiler = [1] * tf.rank(x)
         tiler[-2] = shape[-2]
         out_inv_rep = tf.tile(out_inv, tiler)
 

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 condition.ndim == 2:
+            if tf.rank(condition) == 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 targets.ndim == 3:
+            if len(shape_scale) == 2 and tf.rank(targets) == 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 condition.ndim == 2 else (condition.shape[0], condition.shape[1], 1)
+                (condition.shape[0], 1) if tf.rank(condition) == 2 else (condition.shape[0], condition.shape[1], 1)
             )
             noise_scale = tf.zeros(shape=shape_scale) + 2.0 * self.soft_low