fix(nets): improve VAE reparameterize memory efficiency and revert stray safeeval changes

- Avoid allocating a tensor of ones at inference by returning mu + std
  directly when use_mean_at_inference is disabled, and mu + eps * std
  only during training.
- Revert unintended monai/utils/safeeval.py and tests/utils/test_safe_eval.py
  changes that were accidentally included in this branch.

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

Author: Shizoqua

monai/networks/nets/fullyconnectednet.py

@@ -195,8 +195,10 @@ def reparameterize(self, mu: torch.Tensor, logvar: torch.Tensor) -> torch.Tensor
             # term is only added during training (the reparameterization trick).
             return mu
         std = torch.exp(0.5 * logvar)
-        eps = torch.randn_like(std) if self.training else torch.ones_like(std)
-        return mu + eps * std
+        if self.training:
+            eps = torch.randn_like(std)
+            return mu + eps * std
+        return mu + std
 
     def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
         mu, logvar = self.encode_forward(x)

monai/networks/nets/varautoencoder.py

@@ -165,8 +165,10 @@ def reparameterize(self, mu: torch.Tensor, logvar: torch.Tensor) -> torch.Tensor
             # term is only added during training (the reparameterization trick).
             return mu
         std = torch.exp(0.5 * logvar)
-        eps = torch.randn_like(std) if self.training else torch.ones_like(std)
-        return mu + eps * std
+        if self.training:
+            eps = torch.randn_like(std)
+            return mu + eps * std
+        return mu + std
 
     def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
         mu, logvar = self.encode_forward(x)

monai/utils/safeeval.py

@@ -47,11 +47,12 @@ def __init__(self, int_type_str: str, float_type_str: str):
         self.float_type_str = float_type_str
 
     def visit_Constant(self, node):
-        if isinstance(node.value, (int, float)):
-            type_str = self.int_type_str if isinstance(node.value, int) else self.float_type_str
-            return ast.parse(f"{type_str}({node.value})")
-
-        return node
+        if isinstance(node.value, bool) or not isinstance(node.value, (int, float)):
+            return node
+        type_str = self.int_type_str if isinstance(node.value, int) else self.float_type_str
+        func_node = ast.parse(type_str, mode="eval").body
+        call_node = ast.Call(func=func_node, args=[ast.Constant(value=node.value)], keywords=[])
+        return ast.copy_location(call_node, node)
 
 
 def safe_eval(
@@ -74,7 +75,7 @@ def safe_eval(
     by `int_type_str` and `float_type_str`. These are expected to be constructor names prefixed with `np.` as Numpy
     will be present in the expression global variables under that name. The values can be changed to other types if
     needed, such as "int64". One advantage of doing this is to avoid denial-of-service attacks by attempting to evaluate
-    an expressoini which is incredibly slow under native Python but fast (though potentially erroneous) under Numpy.
+    an expression which is incredibly slow under native Python but fast (though potentially erroneous) under Numpy.
 
     Args:
         expr: expression to evaluate, this will be stripped before parsing to avoid indentation complaints
@@ -101,6 +102,7 @@ def safe_eval(
 
     if rewrite_np:
         parsed = _RewriteConstNp(int_type_str, float_type_str).visit(parsed)
-        locals_vars = {"np": np, **(locals_vars or {})}
+        ast.fix_missing_locations(parsed)
+        locals_vars = {**(locals_vars or {}), "np": np}
 
-    return eval(expr, dict(globals_vars) if globals_vars else None, locals_vars)
+    return eval(compile(parsed, "<safe_eval>", "eval"), dict(globals_vars) if globals_vars else None, locals_vars)

tests/utils/test_safe_eval.py

@@ -14,6 +14,7 @@
 import ast
 import unittest
 
+import numpy as np
 from parameterized import parameterized
 
 from monai.utils import safe_eval
@@ -62,6 +63,35 @@ def test_allowed_types(self):
         with self.assertRaises(ValueError):
             safe_eval("1*2", allowed_types=allowed)
 
+    def test_rewrite_np_produces_numpy_types(self):
+        """Test that rewrite_np wraps literals in numpy types."""
+        result = safe_eval("2 + 3", rewrite_np=True)
+        self.assertIsInstance(result, np.integer)
+
+        result = safe_eval("2.5 + 1.5", rewrite_np=True)
+        self.assertIsInstance(result, np.floating)
+
+    def test_rewrite_np_large_exponent(self):
+        """Test that rewrite_np prevents slow native-Python exponentiation."""
+        # Under native Python, 9**9**9 produces a ~369-million-digit integer;
+        # under np.int32 it overflows and completes almost instantly.
+        result = safe_eval("9**9**9", rewrite_np=True)
+        self.assertIsInstance(result, np.integer)
+
+    def test_rewrite_np_preserves_bool(self):
+        """Test that rewrite_np does not wrap bool constants."""
+        result = safe_eval("True", rewrite_np=True)
+        self.assertIs(result, True)
+
+        result = safe_eval("False", rewrite_np=True)
+        self.assertIs(result, False)
+
+    def test_rewrite_np_inf_constant(self):
+        """Test that rewrite_np handles overflowing infinity literals."""
+        result = safe_eval("1e309", rewrite_np=True)
+        self.assertIsInstance(result, np.floating)
+        self.assertTrue(np.isinf(result))
+
 
 if __name__ == "__main__":
     unittest.main()