fix subnet input dimensions

Author: arrjon

bayesflow/networks/consistency_models/consistency_model.py

@@ -4,7 +4,7 @@
 import numpy as np
 
 from bayesflow.types import Tensor
-from bayesflow.utils import find_network, layer_kwargs, weighted_mean, tensor_utils
+from bayesflow.utils import find_network, layer_kwargs, weighted_mean, tensor_utils, expand_right_as
 from bayesflow.utils.serialization import deserialize, serializable, serialize
 
 from ..inference_network import InferenceNetwork
@@ -343,8 +343,8 @@ def compute_metrics(
 
         log_p = ops.log(p)
         times = keras.random.categorical(ops.expand_dims(log_p, 0), ops.shape(x)[0], seed=self.seed_generator)[0]
-        t1 = ops.take(discretized_time, times)[..., None]
-        t2 = ops.take(discretized_time, times + 1)[..., None]
+        t1 = expand_right_as(ops.take(discretized_time, times), x)
+        t2 = expand_right_as(ops.take(discretized_time, times + 1), x)
 
         # generate noise vector
         noise = keras.random.normal(keras.ops.shape(x), dtype=keras.ops.dtype(x), seed=self.seed_generator)

bayesflow/networks/flow_matching/flow_matching.py

@@ -179,15 +179,17 @@ def _subnet_input(
             The concatenated input tensor for the subnet or a tuple of tensors if concatenation is disabled.
         """
         if self._concatenate_subnet_input:
+            t = keras.ops.broadcast_to(t, keras.ops.shape(x)[:-1] + (1,))
             xtc = tensor_utils.concatenate_valid([x, t, conditions], axis=-1)
             return self.subnet(xtc, training=training)
         else:
+            if training is False:
+                t = keras.ops.broadcast_to(t, keras.ops.shape(x)[:-1] + (1,))
             return self.subnet(x, t, conditions, training=training)
 
     def velocity(self, xz: Tensor, time: float | Tensor, conditions: Tensor = None, training: bool = False) -> Tensor:
         time = keras.ops.convert_to_tensor(time, dtype=keras.ops.dtype(xz))
         time = expand_right_as(time, xz)
-        time = keras.ops.broadcast_to(time, keras.ops.shape(xz)[:-1] + (1,))
 
         subnet_out = self._subnet_input(xz, time, conditions, training=training)
         return self.output_projector(subnet_out, training=training)