Stabilize SO(3) chordal smoother weights across backends

src/pyrecest/smoothers/so3_chordal_mean_smoother.py

@@ -5,9 +5,10 @@
 from operator import index as operator_index
 from typing import Sequence
 
+import numpy as np
+
 # pylint: disable=no-member
-from pyrecest import backend
-from pyrecest.backend import asarray, diag, isfinite, linalg, ndim, zeros
+from pyrecest.backend import asarray, diag, linalg, ndim, to_numpy, zeros
 
 from .abstract_smoother import AbstractSmoother
 
@@ -101,20 +102,32 @@ def _normalize_weight_vector(
             raise ValueError(f"{name} must be one-dimensional.")
         if weights_array.shape[0] != length:
             raise ValueError(f"{name} must have length {length}.")
-        if not bool(backend.all(isfinite(weights_array))):
-            raise ValueError(f"{name} must contain only finite values.")
 
-        for idx in range(length):
-            if weights_array[idx] < 0.0:
-                raise ValueError(f"{name} must be nonnegative.")
+        # Validate and scale using the stored host values. Some XLA operations
+        # flush subnormal operands, and division by a near-maximum finite value
+        # can be lowered through a reciprocal that underflows to zero. Moving
+        # only the resulting O(1) ratios back to the active backend avoids both
+        # failure modes while preserving the intended relative weights.
+        try:
+            host_weights = np.asarray(to_numpy(weights_array))
+            finite_weights = np.isfinite(host_weights)
+        except (TypeError, ValueError, OverflowError, RuntimeError) as exc:
+            raise ValueError(f"{name} must contain real numeric values.") from exc
+        if np.iscomplexobj(host_weights):
+            raise ValueError(f"{name} must contain real numeric values.")
+        if not np.all(finite_weights):
+            raise ValueError(f"{name} must contain only finite values.")
+        if np.any(host_weights < 0.0):
+            raise ValueError(f"{name} must be nonnegative.")
 
-        weight_scale = backend.max(weights_array)
+        weight_scale = float(np.max(host_weights))
         if weight_scale <= 0.0:
             raise ValueError(f"{name} must contain at least one positive entry.")
 
-        scaled_weights = weights_array / weight_scale
+        scaled_host_weights = host_weights / weight_scale
+        scaled_weights = asarray(scaled_host_weights)
         if normalize:
-            return scaled_weights / backend.sum(scaled_weights)
+            return scaled_weights / float(np.sum(scaled_host_weights))
         return scaled_weights
 
     @staticmethod

tests/smoothers/test_so3_chordal_mean_backend_weight_stability.py

@@ -0,0 +1,81 @@
+"""Cross-backend regressions for extreme SO(3) chordal smoother weights."""
+
+import numpy as np
+import numpy.testing as npt
+import pytest
+from pyrecest.backend import array, cos, eye, sin, to_numpy
+from pyrecest.smoothers import SO3ChordalMeanSmoother
+
+# pylint: disable=protected-access
+
+
+def _active_dtype():
+    return to_numpy(array([0.0], dtype=float)).dtype
+
+
+def _z_rotation(angle):
+    return array(
+        [
+            [cos(angle), -sin(angle), 0.0],
+            [sin(angle), cos(angle), 0.0],
+            [0.0, 0.0, 1.0],
+        ]
+    )
+
+
+def test_chordal_mean_preserves_largest_finite_weight_ratio_across_backends():
+    dtype = _active_dtype()
+    largest = np.finfo(dtype).max
+    weights = array(np.asarray([largest, largest / 2.0], dtype=dtype), dtype=float)
+
+    mean_rotation = SO3ChordalMeanSmoother.chordal_mean(
+        [eye(3), _z_rotation(0.5 * np.pi)],
+        weights=weights,
+    )
+
+    npt.assert_allclose(
+        to_numpy(mean_rotation),
+        to_numpy(_z_rotation(np.arctan(0.5))),
+        atol=1.0e-6,
+    )
+
+
+def test_weight_scaling_preserves_positive_subnormal_ratio_across_backends():
+    dtype = _active_dtype()
+    smallest = np.finfo(dtype).smallest_subnormal
+    weights = array(
+        np.asarray([2.0 * smallest, smallest], dtype=dtype),
+        dtype=float,
+    )
+
+    scaled_weights = SO3ChordalMeanSmoother._normalize_weight_vector(
+        weights,
+        2,
+        "weights",
+        normalize=False,
+    )
+    normalized_weights = SO3ChordalMeanSmoother._normalize_weight_vector(
+        weights,
+        2,
+        "weights",
+        normalize=True,
+    )
+
+    npt.assert_allclose(to_numpy(scaled_weights), [1.0, 0.5])
+    npt.assert_allclose(to_numpy(normalized_weights), [2.0 / 3.0, 1.0 / 3.0])
+
+
+def test_rejects_negative_subnormal_weight_across_backends():
+    dtype = _active_dtype()
+    smallest = np.finfo(dtype).smallest_subnormal
+    weights = array(
+        np.asarray([1.0, -smallest], dtype=dtype),
+        dtype=float,
+    )
+
+    with pytest.raises(ValueError, match="nonnegative"):
+        SO3ChordalMeanSmoother._normalize_weight_vector(
+            weights,
+            2,
+            "weights",
+        )