Add optional autograph functionality to all training modes

Author: stefanradev93

bayesflow/coupling_networks.py

@@ -85,6 +85,7 @@ def __init__(self, meta):
         # Optional permutation
         if meta['use_permutation']:
             self.permutation = Permutation(self.latent_dim)
+            self.permutation.trainable = False
         else:
             self.permutation = None
 

bayesflow/helper_functions.py

@@ -20,6 +20,7 @@
 
 import copy
 
+import tensorflow as tf
 from tensorflow.keras.optimizers.schedules import LearningRateSchedule
 
 from bayesflow import default_settings
@@ -104,7 +105,7 @@ def extract_current_lr(optimizer):
 
 def format_loss_string(ep, it, loss, avg_dict, slope=None, lr=None, 
                        ep_str="Epoch", it_str='Iter', scalar_loss_str='Loss'):
-    """ Prepare loss string for displaying on progress bar."""
+    """Prepare loss string for displaying on progress bar."""
 
     disp_str = f"{ep_str}: {ep}, {it_str}: {it}"
     if type(loss) is dict:
@@ -123,6 +124,50 @@ def format_loss_string(ep, it, loss, avg_dict, slope=None, lr=None,
     return disp_str
 
 
+def backprop_step(input_dict, amortizer, optimizer, **kwargs):
+    """Computes the loss of the provided amortizer given an input dictionary and applies gradients.
+
+    Parameters
+    ----------
+    input_dict  : dict
+        The configured output of the genrative model
+    amortizer   : tf.keras.Model
+        The custom amortizer. Needs to implement a compute_loss method.
+    optimizer   : tf.keras.optimizers.Optimizer
+        The optimizer used to update the amortizer's parameters.
+    **kwargs    : dict
+        Optional keyword arguments passed to the network's compute_loss method
+        
+    Returns
+    -------
+    loss : dict
+        The outputs of the compute_loss() method of the amortizer comprising all
+        loss components, such as divergences or regularization.
+    """
+
+    # Forward pass and loss computation
+    with tf.GradientTape() as tape:
+        # Compute custom loss
+        loss = amortizer.compute_loss(input_dict, training=True, **kwargs)
+        # If dict, add components
+        if type(loss) is dict:
+            _loss = tf.add_n(list(loss.values()))
+        else:
+            _loss = loss
+        # Collect regularization loss, if any
+        if amortizer.losses != []:
+            reg = tf.add_n(amortizer.losses)
+            _loss += reg
+            if type(loss) is dict:
+                loss['W.Decay'] = reg
+            else:
+                loss = {'Loss': loss, 'W.Decay': reg}
+    # One step backprop and return loss
+    gradients = tape.gradient(_loss, amortizer.trainable_variables)
+    optimizer.apply_gradients(zip(gradients, amortizer.trainable_variables))
+    return loss
+
+
 def check_posterior_prior_shapes(post_samples, prior_samples):
     """Checks requirements for the shapes of posterior and prior draws as
     necessitated by most diagnostic functions.

bayesflow/helper_networks.py

@@ -118,19 +118,29 @@ def call(self, target, inverse=False):
         ----------
         target   : tf.Tensor of shape (batch_size, ...)
             The target vector to be permuted over its last axis.
-        inverse  : bool, default: False
+        inverse  : bool, optional, default: False
             Controls if the current pass is forward (``inverse=False``) or inverse (``inverse=True``).
 
         Returns
         -------
         out      : tf.Tensor of the same shape as `target`.
-            The permuted target vector.
-
+            The (un-)permuted target vector.
         """
 
         if not inverse:
-            return tf.transpose(tf.gather(tf.transpose(target), self.permutation))
-        return tf.transpose(tf.gather(tf.transpose(target), self.inv_permutation))
+            return self._forward(target)
+        else:
+            return self._inverse(target)
+
+    @tf.function
+    def _forward(self, target):
+        """Performs a fixed permutation over the last axis."""
+        return tf.gather(target, self.permutation, axis=-1)
+    
+    @tf.function
+    def _inverse(self, target):
+        """Un-does the fixed permutation over the last axis."""
+        return tf.gather(target, self.inv_permutation, axis=-1)
 
 
 class ActNorm(tf.keras.Model):