docs(networks): add docstrings to VAE reparameterize methods and tests

Signed-off-by: Lanre Shittu <136805224+Shizoqua@users.noreply.github.com>

Author: Shizoqua

monai/networks/nets/fullyconnectednet.py

@@ -172,6 +172,18 @@ def decode_forward(self, z: torch.Tensor, use_sigmoid: bool = True) -> torch.Ten
         return x
 
     def reparameterize(self, mu: torch.Tensor, logvar: torch.Tensor) -> torch.Tensor:
+        """Sample a latent code using the reparameterization trick.
+
+        At inference (eval mode) the posterior mean is returned directly. During
+        training, returns ``mu + eps * std`` with ``eps ~ N(0, I)``.
+
+        Args:
+            mu: Posterior mean, shape ``(batch, latent_size)``.
+            logvar: Log-variance of the posterior, same shape as ``mu``.
+
+        Returns:
+            Sampled latent code, same shape as ``mu``.
+        """
         if not self.training:
             # At inference the latent code is the posterior mean; the random
             # term is only added during training (the reparameterization trick).

monai/networks/nets/varautoencoder.py

@@ -142,6 +142,18 @@ def decode_forward(self, z: torch.Tensor, use_sigmoid: bool = True) -> torch.Ten
         return x
 
     def reparameterize(self, mu: torch.Tensor, logvar: torch.Tensor) -> torch.Tensor:
+        """Sample a latent code using the reparameterization trick.
+
+        At inference (eval mode) the posterior mean is returned directly. During
+        training, returns ``mu + eps * std`` with ``eps ~ N(0, I)``.
+
+        Args:
+            mu: Posterior mean, shape ``(batch, latent_size)``.
+            logvar: Log-variance of the posterior, same shape as ``mu``.
+
+        Returns:
+            Sampled latent code, same shape as ``mu``.
+        """
         if not self.training:
             # At inference the latent code is the posterior mean; the random
             # term is only added during training (the reparameterization trick).

tests/networks/nets/test_fullyconnectednet.py

@@ -65,8 +65,11 @@ def test_vfc_shape(self, input_param, input_shape, expected_shape):
             self.assertEqual(result.shape, expected_shape)
 
     def test_vfc_reparameterize_eval_returns_mu(self):
-        # At eval the latent code must equal mu (deterministic); at train it must
-        # be stochastic. Same #8413 reparameterize bug as VarAutoEncoder.
+        """A VFC latent code is deterministic at eval (equals mu) and stochastic at train.
+
+        Regression test for the #8413 reparameterize bug, which returned ``mu + std``
+        at inference instead of ``mu``.
+        """
         net = VarFullyConnectedNet(
             in_channels=10, out_channels=10, latent_size=30, encode_channels=(15, 20, 25), decode_channels=(15, 20, 25)
         ).to(device)

tests/networks/nets/test_varautoencoder.py

@@ -123,9 +123,10 @@ def test_script(self):
         test_script_save(net, test_data, rtol=1e-3, atol=1e-3)
 
     def test_reparameterize_eval_returns_mu(self):
-        # At eval the latent code must equal mu (deterministic, no noise added);
-        # at train it must be stochastic. Regression test for #8413, where eval
-        # returned mu + std.
+        """A VarAutoEncoder latent code is deterministic at eval (equals mu) and stochastic at train.
+
+        Regression test for #8413, where eval returned ``mu + std`` instead of ``mu``.
+        """
         net = VarAutoEncoder(
             spatial_dims=2, in_shape=(1, 32, 32), out_channels=1, latent_size=4, channels=(4, 8), strides=(2, 2)
         ).to(device)