Made in-script documentation of model.py more readable

model.py

@@ -160,13 +160,19 @@ class VAE(tf.keras.Model):
     Get a "song map" and make a forward pass through the encoder, in order
     to return the latent representation and the distribution's parameters.
 
-    Parameters:
-    x_input (tf.Tensor): Song map to be encoded by the VAE.
-
-    Returns:
-    tf.Tensor: The parameters of the distribution which encode the song
-               (mu, sd) and a sampled latent representation from this
-               distribution (z_sample).
+    Parameters
+    ----------
+    x_input : tf.Tensor
+        Song map to be encoded by the VAE.
+
+    Returns
+    -------
+    z_sample: tf.Tensor
+        A sampled latent representation from the distribution which encodes the song.
+    mu: tf.Tensor
+        The mean parameter of the distribution.
+    sd: tf.Tensor
+        The standard deviation parameter of the distribution.
     """
       
     mu, rho = tf.split(self.encoder(x_input), num_or_size_splits=2, axis=1)
@@ -178,15 +184,21 @@ class VAE(tf.keras.Model):
     """
     Decode a latent representation of a song.
 
-    Parameters:
-    z_sample (tf.Tensor): Song encoding outputed by the encoder. If
-                          None, this sampling is done over an
-                          unit Gaussian distribution.
+    Parameters
+    ----------
+    z_sample : tf.Tensor
+    
+    Song encoding outputed by the encoder.
+    Default ``None``, for which the sampling is done over an unit Gaussian distribution.
 
-    Returns:
-    tf.Tensor: Song map corresponding to the encoding.
+    Returns
+    -------
+    song_map: tf.Tensor
+        Song map corresponding to the encoding.
     """
       
     if z_sample == None:
-      z_sample = tf.expand_dims(tf.random.normal(shape=(120,)), axis=0)
-    return self.decoder(z_sample)
+      z_sample = tf.expand_dims(tf.random.normal(shape=(120,)), axis=0
+                               
+    song_map = self.decoder(z_sample)
+    return song_map