dims to tuple

Author: arrjon

bayesflow/distributions/diagonal_normal.py

@@ -65,7 +65,7 @@ def build(self, input_shape: Shape) -> None:
         if self.built:
             return
 
-        self.dims = input_shape[1:]
+        self.dims = tuple(input_shape[1:])
 
         self.mean = ops.cast(ops.broadcast_to(self.mean, self.dims), "float32")
         self.std = ops.cast(ops.broadcast_to(self.std, self.dims), "float32")
@@ -91,9 +91,7 @@ def log_prob(self, samples: Tensor, *, normalize: bool = True) -> Tensor:
         result = -0.5 * ops.sum((samples - self._mean) ** 2 / self._std**2, axis=-1)
 
         if normalize:
-            log_normalization_constant = -0.5 * np.sum(self.dims) * math.log(2.0 * math.pi) - ops.sum(
-                ops.log(self._std)
-            )
+            log_normalization_constant = -0.5 * sum(self.dims) * math.log(2.0 * math.pi) - ops.sum(ops.log(self._std))
             result += log_normalization_constant
 
         return result