fix_reverse_out_bound_quadratic_spline (#3140)

* fix_reverse_out_bound_quadratic_spline * fix_reverse_out_bound_quadratic_spline * add two spaces before comment * add sapce * space * lint Co-authored-by: Fritz Obermeyer <[email protected]>
pyro-ppl · Oct 21, 2022 · c573af1 · c573af1
1 parent 0fe132e
commit c573af1
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/pyro/distributions/transforms/spline.py b/pyro/distributions/transforms/spline.py
@@ -254,6 +254,8 @@ def _monotonic_rational_spline(
             c = -input_delta * (inputs - input_cumheights)
 
             discriminant = b.pow(2) - 4 * a * c
+            # Make sure outside_interval input can be reversed as identity.
+            discriminant = discriminant.masked_fill(outside_interval_mask, 0)
             assert (discriminant >= 0).all()
 
             root = (2 * c) / (-b - torch.sqrt(discriminant))