Fix VAE sampling: enable reparameterization by default, return mu for inference

docs/cli.md

@@ -48,6 +48,7 @@ embkit model train-vae INPUT_PATH [OPTIONS]
 | `--seed` | `42` | Random seed |
 | `--bfloat16` | false | Use bfloat16 dtype for reduced memory usage |
 | `--save-stats` | false | Save training statistics alongside the model |
+| `--sampling/--no-sampling` | `true` | Enable reparameterization sampling (VAE). Use `--no-sampling` for a standard autoencoder. |
 
 **Examples**
 

src/embkit/commands/model.py

@@ -38,6 +38,7 @@
 @click.option("--zero-mask", default=None, type=float)
 @click.option("--seed", default=42, type=int)
 @click.option("--bfloat16", is_flag=True)
+@click.option("--sampling/--no-sampling", default=True, show_default=True, help="Enable reparameterization sampling during training (VAE). Use --no-sampling for a standard autoencoder.")
 def train_vae(input_path: str,
               group: str,
               latent: int,
@@ -52,7 +53,8 @@ def train_vae(input_path: str,
               zero_mask: float,
               save_stats: bool,
               seed: int,
-              bfloat16: bool
+              bfloat16: bool,
+              sampling: bool
               ):
     """
     Train VAE model from a TSV file.
@@ -110,6 +112,7 @@ def train_vae(input_path: str,
               latent_dim=latent,
               encoder_layers=enc_layers_list,
               decoder_layers=dec_layers_list,
+              sampling=sampling,
               device=device, dtype=dtype)
 
     loss_func = bce_with_logits
@@ -228,6 +231,6 @@ def encode(input_path: str, model_path:str, normalize:str, out:str):
     m.to(get_device())
     result = m.encoder(df_tensor)
 
-    martix = result[2].detach().cpu().numpy()
-    out_df = pd.DataFrame(martix, index=df.index)
+    matrix = result[0].detach().cpu().numpy()
+    out_df = pd.DataFrame(matrix, index=df.index)
     out_df.to_csv(out, sep="\t")

src/embkit/models/vae/base_vae.py

@@ -39,12 +39,14 @@ def __init__(self, features: List[str], encoder: Optional[Encoder] = None, decod
     def build_encoder(feature_dim: int, latent_dim: int,
                       layers: Optional[LayerList] = None,
                       batch_norm: bool = False,
+                      sampling: bool = True,
                       device=None, dtype=None) -> Encoder:
         return Encoder(
             feature_dim=feature_dim,
             latent_dim=latent_dim,
             layers=layers,
             batch_norm=batch_norm,
+            sampling=sampling,
             device=device, dtype=dtype
         )
 
@@ -73,11 +75,11 @@ def forward(self, x: torch.Tensor):
 
     def encode(self, x:torch.Tensor):
         """
-        Run encoder model and return encoded values
+        Run encoder model and return the latent mean (mu) for stable embeddings.
         """
         with torch.no_grad():
-            _, _, z = self.encoder(x)
-        return z
+            mu, _, _ = self.encoder(x)
+        return mu
 
     @abstractmethod
     def to_dict(self):
@@ -100,8 +102,8 @@ def __init__(self, encoder) -> None:
         self.encoder = encoder
 
     def forward(self, x):
-        _, _, z = self.encoder(x)
-        return z
+        mu, _, _ = self.encoder(x)
+        return mu
 
 def _import_obj(dotted: str):
     """Import 'pkg.mod.ClassName' -> object."""

src/embkit/models/vae/encoder.py

@@ -34,7 +34,7 @@ def __init__(self,
                  batch_norm: bool = False,
                  default_activation: Union[str, None] = "relu",
                  make_latent_heads: bool = True,
-                 sampling : bool = False,
+                 sampling : bool = True,
                  constraint: Optional["NetworkConstraint"] = None,
                  device=None, dtype=None):
         super().__init__()
@@ -113,12 +113,13 @@ def forward(self, x: torch.Tensor):
         if self._make_latent_heads and (self.z_mean is not None) and (self.z_log_var is not None):
             mu = self.z_mean(h)
             logvar = self.z_log_var(h)
-            if self._sampling:
+            if self._sampling and self.training:
                 std = torch.exp(0.5 * logvar)
                 eps = torch.randn_like(std)
                 z = mu + eps * std
-                return mu, logvar, z
-            return mu, logvar, h
+            else:
+                z = mu
+            return mu, logvar, z
 
         return h
 

src/embkit/models/vae/vae.py

@@ -25,6 +25,7 @@ def __init__(
             encoder_layers: Optional[LayerList] = None,
             decoder_layers: Optional[LayerList] = None,
             batch_norm: bool = False,
+            sampling: bool = True,
             device: Optional[torch.device] = None,
             dtype: Optional[torch.dtype] = None
     ):
@@ -35,6 +36,8 @@ def __init__(
             encoder_layers: list of layer configs for Encoder
             decoder_layers: list of layer configs for Decoder
             batch_norm: enable encoder batch normalization blocks
+            sampling: enable reparameterization sampling during the forward pass (VAE).
+                      Set to False for a deterministic autoencoder (mu is used as z).
             device: torch device used for module initialization
             dtype: torch dtype used for module initialization
         """
@@ -52,6 +55,7 @@ def __init__(
         self._encoder_layers_cfg = encoder_layers
         self._decoder_layers_cfg = decoder_layers
         self._batch_norm = batch_norm
+        self._sampling = sampling
         self.latent_dim = latent_dim
 
         feature_dim = len(features)
@@ -63,6 +67,7 @@ def __init__(
             latent_dim=latent_dim,
             layers=encoder_layers,
             batch_norm=batch_norm,
+            sampling=sampling,
             device=device, dtype=dtype
         )
         self.decoder = self.build_decoder(
@@ -87,7 +92,8 @@ def to_dict(self):
             "latent_dim": self.latent_dim,
             "encoder_layers": self._layers_to_dict(self._encoder_layers_cfg),
             "decoder_layers": self._layers_to_dict(self._decoder_layers_cfg),
-            "batch_norm": self._batch_norm
+            "batch_norm": self._batch_norm,
+            "sampling": self._sampling,
         }
 
     @classmethod
@@ -97,6 +103,7 @@ def from_dict(cls, desc):
             latent_dim=desc["latent_dim"],
             encoder_layers=LayerList([Layer.from_dict(li) for li in (desc.get("encoder_layers") or [])]),
             decoder_layers=LayerList([Layer.from_dict(li) for li in (desc.get("decoder_layers") or [])]),
-            batch_norm=desc.get("batch_norm", False)
+            batch_norm=desc.get("batch_norm", False),
+            sampling=desc.get("sampling", True),
         )
 

tests/models/vae_models/test_encoder.py

@@ -111,6 +111,50 @@ def test_encoder_layer_batch_norm(self):
         mu, logvar, z = enc(x)
         assert z.shape == (2, 2)
 
+    def test_sampling_enabled_eval_z_same_as_mu(self):
+        """When sampling=True and eval mode is true, z should not differ from mu, sampling is only on during training."""
+        torch.manual_seed(0)
+        enc = Encoder(feature_dim=8, latent_dim=4, layers=None, sampling=True)
+        enc.eval()  # eval mode disables sampling
+        x = torch.randn(16, 8)
+
+        mu, logvar, z = enc(x)
+
+        self.assertEqual(mu.shape, z.shape)
+        # With sampling enabled, z should equal mu (reparameterization is disabled durning eval)
+        self.assertTrue(torch.allclose(z, mu),
+                         "z should not differ from mu when sampling=True and Eval mode is on")
+
+    def test_sampling_disabled_z_equals_mu(self):
+        """When sampling=False, z should equal mu (no reparameterization noise)."""
+        enc = Encoder(feature_dim=8, latent_dim=4, layers=None, sampling=False)
+        x = torch.randn(16, 8)
+
+        mu, logvar, z = enc(x)
+
+        self.assertEqual(mu.shape, z.shape)
+        # With sampling disabled, z must be identical to mu
+        self.assertTrue(torch.allclose(z, mu),
+                        "z should equal mu when sampling=False")
+
+    def test_default_sampling_is_true(self):
+        """Encoder default should have sampling enabled for proper VAE training."""
+        enc = Encoder(feature_dim=6, latent_dim=3)
+        self.assertTrue(enc._sampling, "Default sampling should be True")
+
+    def test_forward_always_returns_three_tuple_with_latent_heads(self):
+        """forward() must always return (mu, logvar, z) when make_latent_heads=True."""
+        for sampling in (True, False):
+            enc = Encoder(feature_dim=6, latent_dim=3, sampling=sampling)
+            x = torch.randn(4, 6)
+            result = enc(x)
+            self.assertIsInstance(result, tuple)
+            self.assertEqual(len(result), 3,
+                             f"Expected 3-tuple with sampling={sampling}")
+            mu, logvar, z = result
+            self.assertEqual(mu.shape, z.shape)
+            self.assertEqual(mu.shape, logvar.shape)
+
 
 if __name__ == "__main__":
     unittest.main()