Temporary rebase staging for #4294

src/pyrecest/distributions/circle/piecewise_constant_distribution.py

@@ -4,6 +4,7 @@
 import numpy as np
 import pyrecest.backend
 from pyrecest.backend import (
+    amax,
     arange,
     array,
     exp,
@@ -144,11 +145,12 @@ def __init__(self, w):
         if any(bool(weight < 0.0) for weight in w):
             raise ValueError("Weights must be nonnegative")
 
-        mean_weight = mean(w)
-        if not bool(mean_weight > 0.0):
+        weight_scale = amax(w)
+        if not bool(weight_scale > 0.0):
             raise ValueError("Weights must have positive total mass")
+        scaled_weights = w / weight_scale
 
-        self.w = w / (mean_weight * 2.0 * pi)
+        self.w = scaled_weights / (mean(scaled_weights) * 2.0 * pi)
 
     def pdf(self, xs):
         """Evaluate the pdf at each point in xs.

tests/distributions/test_piecewise_constant_weight_overflow.py

@@ -0,0 +1,23 @@
+import numpy as np
+import numpy.testing as npt
+
+from pyrecest.backend import array, to_numpy
+from pyrecest.distributions.circle.piecewise_constant_distribution import (
+    PiecewiseConstantDistribution,
+)
+
+
+def test_maximum_finite_weights_normalize_without_overflow():
+    backend_dtype = np.asarray(to_numpy(array([1.0], dtype=float))).dtype
+    maximum_finite_weight = np.finfo(backend_dtype).max
+
+    distribution = PiecewiseConstantDistribution(
+        array([maximum_finite_weight, maximum_finite_weight / 2.0], dtype=float)
+    )
+
+    npt.assert_allclose(
+        np.asarray(to_numpy(distribution.w)),
+        np.array([2.0 / (3.0 * np.pi), 1.0 / (3.0 * np.pi)]),
+        rtol=1.0e-6,
+        atol=0.0,
+    )