Stabilize tensor-train relative truncation

src/pyrecest/distributions/hypertorus/_tensor_train.py

@@ -65,21 +65,40 @@ def _check_dense_validation_size(size, max_entries):
         )
 
 
+def _stable_vector_norm(values):
+    """Return a scale-safe Euclidean norm for finite array values."""
+    magnitudes = np.abs(np.asarray(values))
+    if magnitudes.size == 0:
+        return 0.0
+    scale = float(np.max(magnitudes))
+    if scale == 0.0 or not np.isfinite(scale):
+        return scale
+    return scale * sqrt(float(np.sum((magnitudes / scale) ** 2)))
+
+
 def _choose_rank(singular_values, max_rank, local_tolerance):
     max_rank = _normalize_max_rank(max_rank)
     full_rank = singular_values.size
     if local_tolerance <= 0:
         rank = full_rank
     else:
-        squared_tail = np.cumsum(singular_values[::-1] ** 2)[::-1]
-        rank = full_rank
-        for candidate in range(1, full_rank + 1):
-            tail = (
-                0.0 if candidate == full_rank else sqrt(float(squared_tail[candidate]))
-            )
-            if tail <= local_tolerance:
-                rank = candidate
-                break
+        scale = float(np.max(np.abs(singular_values)))
+        if scale == 0.0:
+            rank = 1
+        else:
+            scaled_values = singular_values / scale
+            squared_tail = np.cumsum(scaled_values[::-1] ** 2)[::-1]
+            scaled_tolerance = local_tolerance / scale
+            rank = full_rank
+            for candidate in range(1, full_rank + 1):
+                tail = (
+                    0.0
+                    if candidate == full_rank
+                    else sqrt(float(squared_tail[candidate]))
+                )
+                if tail <= scaled_tolerance:
+                    rank = candidate
+                    break
     if max_rank is not None:
         rank = min(rank, max_rank)
     return max(1, rank)
@@ -153,7 +172,7 @@ def from_dense(cls, tensor, *, max_rank=None, rtol=0.0, atol=0.0):
         if array.ndim == 1:
             return cls((array.reshape(1, array.shape[0], 1),))
 
-        norm = float(np.linalg.norm(array.ravel()))
+        norm = _stable_vector_norm(array.ravel())
         global_tolerance = max(atol, rtol * norm)
         local_tolerance = (
             global_tolerance / sqrt(array.ndim - 1) if global_tolerance > 0 else 0.0

tests/distributions/test_hypertoroidal_tensor_train.py

@@ -13,6 +13,14 @@ def test_dense_roundtrip_and_entry_access(self):
         self.assertEqual(tt.shape, (3, 3, 3))
         npt.assert_allclose(tt.entry((1, 2, 0)), tensor[1, 2, 0], atol=1e-12)
 
+    def test_relative_truncation_handles_large_finite_values(self):
+        tensor = np.diag([1e200, 1e200])
+        with np.errstate(over="raise", invalid="raise"):
+            tt = TensorTrain.from_dense(tensor, rtol=1e-12)
+
+        self.assertEqual(tt.ranks, (1, 2, 1))
+        npt.assert_allclose(tt.to_dense() / 1e200, np.eye(2), atol=1e-12)
+
     def test_frobenius_norm(self):
         tensor = np.arange(9, dtype=float).reshape(3, 3) - 2.0
         tt = TensorTrain.from_dense(tensor)