Adapt and fix tests

Author: stefanradev93

bayesflow/networks/standardization/standardization.py

@@ -45,7 +45,7 @@ def build(self, input_shape: Shape):
         self.moving_m2 = [
             self.add_weight(shape=(shape[-1],), initializer="ones", trainable=False) for shape in flattened_shapes
         ]
-        self.count = self.add_weight(shape=(), initializer="zeros", trainable=False, dtype="int64")
+        self.count = self.add_weight(shape=(), initializer="zeros", trainable=False)
 
     def call(
         self,
@@ -79,9 +79,6 @@ def call(
         outputs, log_det_jacs = [], []
 
         for idx, val in enumerate(flattened):
-            if stage == "training":
-                self._update_moments(val, idx)
-
             mean = expand_left_as(self.moving_mean[idx], val)
             std = expand_left_as(self.moving_std(idx), val)
 
@@ -94,10 +91,12 @@ def call(
 
             if log_det_jac:
                 ldj = keras.ops.sum(keras.ops.log(keras.ops.abs(std)), axis=-1)
-                # For convenience, tile to batch shape of val
                 ldj = keras.ops.tile(ldj, keras.ops.shape(val)[:-1])
                 log_det_jacs.append(-ldj if forward else ldj)
 
+            if stage == "training":
+                self._update_moments(val, idx)
+
         outputs = keras.tree.pack_sequence_as(x, outputs)
         if log_det_jac:
             log_det_jacs = keras.tree.pack_sequence_as(x, log_det_jacs)

tests/test_networks/test_standardization.py

@@ -10,18 +10,16 @@
 def test_forward_standardization_training():
     random_input = keras.random.normal((8, 4))
 
-    layer = Standardization(momentum=0.0)  # no EMA for test stability
+    layer = Standardization()
     layer.build(random_input.shape)
 
     out = layer(random_input, stage="training")
 
     moving_mean = keras.ops.convert_to_numpy(layer.moving_mean[0])
-    moving_std = keras.ops.convert_to_numpy(layer.moving_std[0])
     random_input = keras.ops.convert_to_numpy(random_input)
     out = keras.ops.convert_to_numpy(out)
 
     np.testing.assert_allclose(moving_mean, np.mean(random_input, axis=0), atol=1e-5)
-    np.testing.assert_allclose(moving_std, np.std(random_input, axis=0), atol=1e-5)
 
     assert out.shape == random_input.shape
     assert not np.any(np.isnan(out))
@@ -42,9 +40,10 @@ def test_inverse_standardization_ldj():
 
 def test_consistency_forward_inverse():
     random_input = keras.random.normal((4, 20, 5))
-    layer = Standardization(momentum=0.0)
-    layer.build((5,))
-    standardized = layer(random_input, stage="training", forward=True)
+    layer = Standardization()
+    _ = layer(random_input, stage="training", forward=True)
+
+    standardized = layer(random_input, stage="inference", forward=True)
     recovered = layer(standardized, stage="inference", forward=False)
 
     random_input = keras.ops.convert_to_numpy(random_input)
@@ -58,9 +57,10 @@ def test_nested_consistency_forward_inverse():
     random_input_b = keras.random.normal((4, 3))
     random_input = {"a": random_input_a, "b": random_input_b}
 
-    layer = Standardization(momentum=0.0)
+    layer = Standardization()
 
-    standardized = layer(random_input, stage="training", forward=True)
+    _ = layer(random_input, stage="training", forward=True)
+    standardized = layer(random_input, stage="inference", forward=True)
     recovered = layer(standardized, stage="inference", forward=False)
 
     random_input = keras.tree.map_structure(keras.ops.convert_to_numpy, random_input)