fix function name

Author: arrjon

bayesflow/experimental/diffusion_model/diffusion_model.py

@@ -200,7 +200,7 @@ def convert_prediction_to_x(
             return (z + sigma_t**2 * pred) / alpha_t
         raise ValueError(f"Unknown prediction type {self._prediction_type}.")
 
-    def _subnet_input(
+    def _apply_subnet(
         self, xz: Tensor, log_snr: Tensor, conditions: Tensor = None, training: bool = False
     ) -> Tensor | tuple[Tensor, Tensor, Tensor]:
         """
@@ -270,7 +270,7 @@ def velocity(
         log_snr_t = ops.broadcast_to(log_snr_t, ops.shape(xz)[:-1] + (1,))
         alpha_t, sigma_t = self.noise_schedule.get_alpha_sigma(log_snr_t=log_snr_t)
 
-        subnet_out = self._subnet_input(
+        subnet_out = self._apply_subnet(
             xz, self._transform_log_snr(log_snr_t), conditions=conditions, training=training
         )
         pred = self.output_projector(subnet_out, training=training)
@@ -491,7 +491,7 @@ def compute_metrics(
         diffused_x = alpha_t * x + sigma_t * eps_t
 
         # calculate output of the network
-        subnet_out = self._subnet_input(
+        subnet_out = self._apply_subnet(
             diffused_x, self._transform_log_snr(log_snr_t), conditions=conditions, training=training
         )
         pred = self.output_projector(subnet_out, training=training)

bayesflow/networks/consistency_models/consistency_model.py

@@ -258,7 +258,7 @@ def _inverse(self, z: Tensor, conditions: Tensor = None, training: bool = False,
             x = self.consistency_function(x_n, t, conditions=conditions, training=training)
         return x
 
-    def _subnet_input(
+    def _apply_subnet(
         self, x: Tensor, t: Tensor, conditions: Tensor = None, training: bool = False
     ) -> Tensor | tuple[Tensor, Tensor, Tensor]:
         """
@@ -302,7 +302,7 @@ def consistency_function(self, x: Tensor, t: Tensor, conditions: Tensor = None,
             Whether internal layers (e.g., dropout) should behave in train or inference mode.
         """
 
-        subnet_out = self._subnet_input(x, t, conditions, training=training)
+        subnet_out = self._apply_subnet(x, t, conditions, training=training)
         f = self.output_projector(subnet_out)
 
         # Compute skip and out parts (vectorized, since self.sigma2 is of shape (1, input_dim)

bayesflow/networks/flow_matching/flow_matching.py

@@ -155,7 +155,7 @@ def get_config(self):
 
         return base_config | serialize(config)
 
-    def _subnet_input(
+    def _apply_subnet(
         self, x: Tensor, t: Tensor, conditions: Tensor = None, training: bool = False
     ) -> Tensor | tuple[Tensor, Tensor, Tensor]:
         """
@@ -191,7 +191,7 @@ def velocity(self, xz: Tensor, time: float | Tensor, conditions: Tensor = None,
         time = keras.ops.convert_to_tensor(time, dtype=keras.ops.dtype(xz))
         time = expand_right_as(time, xz)
 
-        subnet_out = self._subnet_input(xz, time, conditions, training=training)
+        subnet_out = self._apply_subnet(xz, time, conditions, training=training)
         return self.output_projector(subnet_out, training=training)
 
     def _velocity_trace(