Explicitly specify domain for AffineTransform in ScaledLogitTransform (…

…#378)
CDCgov · Aug 14, 2024 · 85726bd · 85726bd
1 parent 0fc0c8d
commit 85726bd
Showing 1 changed file with 7 additions and 1 deletion.
diff --git a/model/src/pyrenew/transformation/builtin.py b/model/src/pyrenew/transformation/builtin.py
@@ -3,6 +3,7 @@
 """
 
 import numpyro.distributions.transforms as nt
+from numpyro.distributions import constraints
 
 
 def ScaledLogitTransform(
@@ -27,5 +28,10 @@ def ScaledLogitTransform(
         - numpyro.distributions.transforms.SigmoidTransform().inv
     """
     return nt.ComposeTransform(
-        [nt.AffineTransform(0.0, 1.0 / x_max), nt.SigmoidTransform().inv]
+        [
+            nt.AffineTransform(
+                0.0, 1.0 / x_max, domain=constraints.interval(0, x_max)
+            ),
+            nt.SigmoidTransform().inv,
+        ]
     )