Refactor to remove set_head_shapes_from_target_shape

Author: han-ol

bayesflow/networks/point_inference_network.py

@@ -52,12 +52,12 @@ def build(self, xz_shape: Shape, conditions_shape: Shape = None) -> None:
         self.heads_flat = dict()  # see comment regarding heads_flat below
 
         for score_key, score in self.scores.items():
-            score.set_head_shapes_from_target_shape(xz_shape)
+            head_shapes = score.get_head_shapes_from_target_shape(xz_shape)
 
             self.heads[score_key] = {}
 
-            for head_key in score.head_shapes.keys():
-                head = score.get_head(head_key)
+            for head_key, head_shape in head_shapes.items():
+                head = score.get_head(head_key, head_shape)
                 head.build(body_output_shape)
                 # If head is not tracked explicitly, self.variables does not include them.
                 # Testing with tests.utils.assert_layers_equal() would thus neglect heads

bayesflow/scores/scores.py

@@ -60,10 +60,20 @@ def from_config(cls, config):
     def get_head_shapes_from_target_shape(self, target_shape):
         raise NotImplementedError
 
-    def set_head_shapes_from_target_shape(self, target_shape):
-        self.head_shapes = self.get_head_shapes_from_target_shape(target_shape)
-
-    def get_subnet(self, key: str):
+    def get_subnet(self, key: str) -> keras.Layer:
+        """For a specified key, request a subnet to be used for projecting the shared condition embedding
+        before reshaping to the heads output shape.
+
+        Parameters
+        ----------
+        key : str
+            Name of head for which to request a link.
+
+        Returns
+        -------
+        link : keras.Layer
+            Subnet projecting the shared condition embedding.
+        """
         if key not in self.subnets.keys():
             return keras.layers.Identity()
         else:
@@ -77,15 +87,15 @@ def get_link(self, key: str):
         else:
             return self.links[key]
 
-    def get_head(self, key: str):
+    def get_head(self, key: str, shape: Shape):
         subnet = self.get_subnet(key)
-        head_shape = self.head_shapes[key]
-        dense = keras.layers.Dense(units=math.prod(head_shape))
-        reshape = keras.layers.Reshape(target_shape=head_shape)
+        dense = keras.layers.Dense(units=math.prod(shape))
+        reshape = keras.layers.Reshape(target_shape=shape)
         link = self.get_link(key)
         return keras.Sequential([subnet, dense, reshape, link])
 
     def score(self, estimates: dict[str, Tensor], targets: Tensor, weights: Tensor) -> Tensor:
+        """Scores a probabilistic estimate based of a distribution based on samples of that distribution."""
         raise NotImplementedError
 
     def aggregate(self, scores: Tensor, weights: Tensor = None):
@@ -114,7 +124,7 @@ def __init__(
             "k": k,
         }
 
-    def get_head_shapes_from_target_shape(self, target_shape):
+    def get_head_shapes_from_target_shape(self, target_shape: Shape):
         # keras.saving.load_model sometimes passes target_shape as a list.
         # This is why I force a conversion to tuple here.
         target_shape = tuple(target_shape)
@@ -180,7 +190,7 @@ def get_config(self):
         base_config = super().get_config()
         return base_config | self.config
 
-    def get_head_shapes_from_target_shape(self, target_shape):
+    def get_head_shapes_from_target_shape(self, target_shape: Shape):
         # keras.saving.load_model sometimes passes target_shape as a list.
         # This is why I force a conversion to tuple here.
         target_shape = tuple(target_shape)
@@ -240,7 +250,7 @@ def get_config(self):
         base_config = super().get_config()
         return base_config | self.config
 
-    def get_head_shapes_from_target_shape(self, target_shape) -> dict[str, Shape]:
+    def get_head_shapes_from_target_shape(self, target_shape: Shape) -> dict[str, Shape]:
         self.D = target_shape[-1]
         return dict(
             mean=(self.D,),

tests/test_networks/test_point_inference_network/test_point_inference_network.py

@@ -12,9 +12,10 @@ def test_output_structure(point_inference_network, random_samples, random_condit
 
     assert isinstance(output, dict)
     for score_key, score in point_inference_network.scores.items():
-        assert isinstance(score.head_shapes, dict)
+        head_shapes = score.get_head_shapes_from_target_shape(random_samples.shape)
+        assert isinstance(head_shapes, dict)
 
-        for head_key, head_shape in score.head_shapes.items():
+        for head_key, head_shape in head_shapes.items():
             head_output = output[score_key][head_key]
             assert keras.ops.is_tensor(head_output)
             assert head_output.shape[1:] == head_shape

tests/test_scores/test_scores.py

@@ -15,11 +15,12 @@ def test_score_output(scoring_rule, random_conditions):
     if random_conditions is None:
         random_conditions = keras.ops.convert_to_tensor([[1.0]])
 
-    scoring_rule.set_head_shapes_from_target_shape(random_conditions.shape)
+    # Using random random_conditions also as targets for the purpose of this test.
+    head_shapes = scoring_rule.get_head_shapes_from_target_shape(random_conditions.shape)
     print(scoring_rule.get_config())
     estimates = {
         k: scoring_rule.get_link(k)(keras.random.normal((random_conditions.shape[0],) + head_shape))
-        for k, head_shape in scoring_rule.head_shapes.items()
+        for k, head_shape in head_shapes.items()
     }
     score = scoring_rule.score(estimates, random_conditions)
 
@@ -30,7 +31,6 @@ def test_mean_score_optimality(mean_score, random_conditions):
     if random_conditions is None:
         random_conditions = keras.ops.convert_to_tensor([[1.0]])
 
-    mean_score.set_head_shapes_from_target_shape(random_conditions.shape)
     key = "value"
     suboptimal_estimates = {key: keras.random.uniform(random_conditions.shape)}
     optimal_estimates = {key: random_conditions}