fix: coupling flows and CMs not working on GPU due to int type (#273)

Apparently, int32 variables are not transferred to the GPU,
leading to problems with XLA. Changing the type declarations
to int seems to fix the problem

Author: vpratz

bayesflow/networks/consistency_models/consistency_model.py

@@ -83,7 +83,7 @@ def __init__(
         self.s0 = float(s0)
         self.s1 = float(s1)
         # create variable that works with JIT compilation
-        self.current_step = self.add_weight(name="current_step", initializer="zeros", trainable=False, dtype="int32")
+        self.current_step = self.add_weight(name="current_step", initializer="zeros", trainable=False, dtype="int")
         self.current_step.assign(0)
 
         self.seed_generator = keras.random.SeedGenerator()
@@ -258,7 +258,7 @@ def compute_metrics(self, x: Tensor, conditions: Tensor = None, stage: str = "tr
             self.current_step.assign(ops.minimum(self.current_step, self.total_steps - 1))
 
         discretization_index = ops.take(
-            self.discretization_map, ops.cast(self._schedule_discretization(self.current_step), "int32")
+            self.discretization_map, ops.cast(self._schedule_discretization(self.current_step), "int")
         )
         discretized_time = ops.take(self.discretized_times, discretization_index, axis=0)
 

bayesflow/networks/coupling_flow/permutations/random.py

@@ -16,12 +16,12 @@ def build(self, xz_shape: Shape, **kwargs) -> None:
             shape=(xz_shape[-1],),
             initializer=keras.initializers.Constant(forward_indices),
             trainable=False,
-            dtype="int32",
+            dtype="int",
         )
 
         self.inverse_indices = self.add_weight(
             shape=(xz_shape[-1],),
             initializer=keras.initializers.Constant(inverse_indices),
             trainable=False,
-            dtype="int32",
+            dtype="int",
         )

bayesflow/networks/coupling_flow/permutations/swap.py

@@ -16,12 +16,12 @@ def build(self, xz_shape: Shape, **kwargs) -> None:
             shape=(xz_shape[-1],),
             initializer=keras.initializers.Constant(forward_indices),
             trainable=False,
-            dtype="int32",
+            dtype="int",
         )
 
         self.inverse_indices = self.add_variable(
             shape=(xz_shape[-1],),
             initializer=keras.initializers.Constant(inverse_indices),
             trainable=False,
-            dtype="int32",
+            dtype="int",
         )