From a8c0f0a1fa04f86a92ce93d83d16398020aef2ce Mon Sep 17 00:00:00 2001
From: boatdrinks100 <ian.pearman100@gmail.com>
Date: Tue, 10 Mar 2026 13:02:01 +0000
Subject: [PATCH 01/11] small changes to avoid numerical instability

---
 jammy_flows/layers/spheres/sphere_base.py | 8 +++++++-
 1 file changed, 7 insertions(+), 1 deletion(-)

diff --git a/jammy_flows/layers/spheres/sphere_base.py b/jammy_flows/layers/spheres/sphere_base.py
index 18b0ce0..9f46f91 100644
--- a/jammy_flows/layers/spheres/sphere_base.py
+++ b/jammy_flows/layers/spheres/sphere_base.py
@@ -193,7 +193,13 @@ def eucl_to_spherical_embedding(self, x, log_det):
                 else:
 
                     # last one is 0 to 2pi
-                    new_angle=torch.acos(x[:,ind:ind+1]/torch.sum(x[:,ind:]**2, dim=1, keepdims=True).sqrt())
+                    # new_angle=torch.acos(x[:,ind:ind+1]/torch.sum(x[:,ind:]**2, dim=1, keepdims=True).sqrt())
+                    sq = torch.sum(x[:, ind:] ** 2, dim=1, keepdims=True)
+                    denom = torch.sqrt(torch.clamp(sq, min=1e-8))
+                    arg = x[:, ind:ind+1] / denom
+                    arg = torch.clamp(arg, -1.0, 1.0)
+                    new_angle = torch.acos(arg)
+
                     #mask_smaller=(x[:,ind+1:ind+2]<0).double()
                     new_angle=torch.where(x[:,ind+1:ind+2]<0, 2*numpy.pi-new_angle, new_angle)
 

From 2adc94cfa4803238cee253397698eb5be12fa5be Mon Sep 17 00:00:00 2001
From: James Pearman <james@gridraven.com>
Date: Fri, 17 Apr 2026 00:08:55 +0300
Subject: [PATCH 02/11] add compilability to o flow

---
 jammy_flows/layers/spline_fns.py |  9 ++----
 jammy_flows/main/default.py      | 53 +++++++-------------------------
 2 files changed, 14 insertions(+), 48 deletions(-)

diff --git a/jammy_flows/layers/spline_fns.py b/jammy_flows/layers/spline_fns.py
index 307997c..799d943 100644
--- a/jammy_flows/layers/spline_fns.py
+++ b/jammy_flows/layers/spline_fns.py
@@ -11,9 +11,10 @@
     print("Sympy not installed!")
 
 def searchsorted(bin_locations, inputs, eps=1e-6):
-    bin_locations[..., -1] += eps
+    last = bin_locations[..., -1:] + eps
+    adjusted = torch.cat([bin_locations[..., :-1], last], dim=-1)
     return torch.sum(
-        inputs >= bin_locations,
+        inputs >= adjusted,
         dim=-1,
         keepdims=True
     ) - 1
@@ -33,10 +34,6 @@ def rational_quadratic_spline(inputs,
                               restrict_max_min_width_height_ratio=-1.0):
 
         
-        if torch.min(inputs) < left or torch.max(inputs) > right:
-           
-            raise Exception("outside boundaries in rational-spline flow! (min/max (%.2f/%.2f), allowed: (%.2f/%.2f)" % (torch.min(inputs), torch.max(inputs), left, right))
-
         num_bins = unnormalized_widths.shape[-1]
 
         if rel_min_bin_width * num_bins > 1.0:
diff --git a/jammy_flows/main/default.py b/jammy_flows/main/default.py
index ce3f705..fbd0de8 100644
--- a/jammy_flows/main/default.py
+++ b/jammy_flows/main/default.py
@@ -899,8 +899,6 @@ def all_layer_inverse(self,
         extra_conditional_input=[]
         base_targets=[]
 
-        individual_logps=dict()
-
         extra_params = None
 
         if(amortization_parameters is not None):
@@ -992,31 +990,10 @@ def all_layer_inverse(self,
                             - layer.total_param_num : -extra_param_counter,
                         ]
 
-                
-                if(l==(len(pdf_layers)-1)):
-                    # force embedding or intrinsic coordinates in the layer that defines the target dimension
-                    this_target, log_det = layer.inv_flow_mapping([this_target, log_det], extra_inputs=this_extra_params)
-                else:
+                this_target, log_det = layer.inv_flow_mapping([this_target, log_det], extra_inputs=this_extra_params)
 
-                    this_target, log_det = layer.inv_flow_mapping([this_target, log_det], extra_inputs=this_extra_params)
-                
                 extra_param_counter += layer.total_param_num
 
-            if(False):
-                ## stems from debugging purposes, not used currently
-                ind_base_eval=this_logp = torch.distributions.MultivariateNormal(
-                    torch.zeros_like(this_target).to(x),
-                    covariance_matrix=torch.eye(this_target.shape[1]).type_as(x).to(x),
-                ).log_prob(this_target)
-
-                ind_logdet=log_det
-                
-
-                individual_logps["%.2d_%s" % (pdf_index, this_pdf_type)]=ind_base_eval+ind_logdet
-                individual_logps["%.2d_%s_logdet" % (pdf_index, this_pdf_type)]=ind_logdet
-                individual_logps["%.2d_%s_base" % (pdf_index, this_pdf_type)]=ind_base_eval
-
-            
             base_targets.append(this_target)
 
             prev_target=x[:,self.target_dim_indices[pdf_index][0]:self.target_dim_indices[pdf_index][1]]
@@ -1174,12 +1151,11 @@ def obtain_flow_param_structure(self,
         else:
 
             if(seed is not None):
-                numpy.random.seed(seed)
+                torch.manual_seed(seed)
 
-            std_normal = numpy.random.normal(size=(used_sample_size, self.total_base_dim))
-
-            std_normal_samples = (
-                torch.from_numpy(std_normal).type(data_type).to(used_device)
+            std_normal_samples = torch.randn(
+                used_sample_size, self.total_base_dim,
+                dtype=data_type, device=used_device,
             )
             log_gauss_evals = torch.distributions.MultivariateNormal(
                 torch.zeros(self.total_base_dim).type(data_type).to(used_device),
@@ -1454,11 +1430,8 @@ def all_layer_forward(self,
                     
                     this_extra_params = extra_params[:, extra_param_counter : extra_param_counter + layer.total_param_num]
               
-                if(l==(len(pdf_layers)-1)):
-                    this_target, log_det = layer.flow_mapping([this_target, log_det], extra_inputs=this_extra_params)
-                else:
-                    this_target, log_det = layer.flow_mapping([this_target, log_det], extra_inputs=this_extra_params)
-                
+                this_target, log_det = layer.flow_mapping([this_target, log_det], extra_inputs=this_extra_params)
+
                 extra_param_counter += layer.total_param_num
 
             new_targets.append(this_target)
@@ -1469,9 +1442,6 @@ def all_layer_forward(self,
 
             extra_conditional_input.append(prev_target)
 
-        if (torch.isfinite(x) == 0).sum() > 0:
-            raise Exception("nonfinite samples generated .. this should never happen!")
-        
         x=torch.cat(new_targets, dim=1)
 
         ## transform to desired output space 
@@ -1614,12 +1584,11 @@ def _obtain_sample(self,
         else:
 
             if(seed is not None):
-                numpy.random.seed(seed)
-
-            std_normal = numpy.random.normal(size=(used_sample_size, self.total_base_dim))
+                torch.manual_seed(seed)
 
-            std_normal_samples = (
-                torch.from_numpy(std_normal).type(data_type).to(used_device)
+            std_normal_samples = torch.randn(
+                used_sample_size, self.total_base_dim,
+                dtype=data_type, device=used_device,
             )
 
             log_gauss_evals=torch.distributions.Normal(0.0,1.0).log_prob(std_normal_samples).sum(dim=-1)

From bf9526b13ac5949e3aaee3b55710f10af393655b Mon Sep 17 00:00:00 2001
From: James Pearman <james@gridraven.com>
Date: Fri, 17 Apr 2026 12:19:18 +0300
Subject: [PATCH 03/11] Add compilation for g flows, monkey patched into this
 for fp32 support

---
 jammy_flows/layers/bisection_n_newton.py | 93 +++++++-----------------
 jammy_flows/layers/spline_fns.py         | 22 ++++--
 2 files changed, 41 insertions(+), 74 deletions(-)

diff --git a/jammy_flows/layers/bisection_n_newton.py b/jammy_flows/layers/bisection_n_newton.py
index 31b00b6..5617ac8 100644
--- a/jammy_flows/layers/bisection_n_newton.py
+++ b/jammy_flows/layers/bisection_n_newton.py
@@ -38,89 +38,46 @@ def inverse_bisection_n_newton_joint_func_and_grad(func,
             The inverse of the function *func* in each sub-dimension in each batch item.
 
     """
-    new_upper = torch.tensor(max_boundary).type(target_arg.dtype).repeat(*target_arg.shape).to(target_arg.device)
-    new_lower = torch.tensor(min_boundary).type(target_arg.dtype).repeat(*target_arg.shape).to(target_arg.device)
- 
-    mid=0
+    new_upper = torch.full_like(target_arg, max_boundary)
+    new_lower = torch.full_like(target_arg, min_boundary)
+
+    mid = (new_upper + new_lower) / 2.
     for i in range(num_bisection_iter):
         mid = (new_upper + new_lower) / 2.
-        #print("mid: ", mid)
         inverse_mid = func(mid, *args)
 
-        #print("MID", mid)
-        
-        right_part = (inverse_mid < target_arg).type(target_arg.dtype)
+        right_part = (inverse_mid < target_arg).to(target_arg.dtype)
         left_part = 1. - right_part
 
-        correct_part = (close(inverse_mid, target_arg, rtol=1e-6, atol=0)).type(target_arg.dtype)
+        correct_part = close(inverse_mid, target_arg, rtol=1e-6, atol=0).to(target_arg.dtype)
 
         new_lower = (1. - correct_part) * (right_part * mid + left_part * new_lower) + correct_part * mid
         new_upper = (1. - correct_part) * (right_part * new_upper + left_part * mid) + correct_part * mid
-      
-        
-    prev=mid
-
-    #print("target arg", target_arg.shape)
-
-
-    above_tolerance_mask=torch.ones( target_arg.shape[0], dtype=torch.bool, device=target_arg.device)
-
-    ## check where we want to broadcast the masking, and wnhere not
-
-    broadcasting_bool_args=[True if (prev.shape[0]>1 and arg.shape[0]>1) else False for arg in args ]
 
+    prev = mid
+
+    # Newton iterations — compile-friendly variant.  The original code sliced
+    # `prev[above_tolerance_mask, :]` each iter to only operate on unconverged
+    # rows, but boolean-mask indexing produces a tensor with an unbacked symint
+    # size which breaks `torch.compile(fullgraph=True)`.  Instead, compute on
+    # the full batch every iteration and gate the update with `torch.where`
+    # so converged rows stay frozen.  Fixed `num_newton_iter` is done — no
+    # data-dependent early exit.  For our use case (D=1, num_newton_iter=20)
+    # the extra compute on converged rows is negligible and compile-trace
+    # fusion more than makes up for it.
+    active_row = torch.ones(target_arg.shape[0], dtype=torch.bool, device=target_arg.device)
     for i in range(num_newton_iter):
-       
-        fn_result, f_prime_eval = joint_func(prev[above_tolerance_mask,:], *[a[above_tolerance_mask] if(broadcasting_bool_args[arg_index] == True) else a for arg_index, a in enumerate(args)])
-        
-        f_eval=fn_result-target_arg[above_tolerance_mask,:]
-
-        update=(f_eval/f_prime_eval)
-
-        newsource=prev[above_tolerance_mask,:]-update
-
-        prev=torch.masked_scatter(input=prev, mask=above_tolerance_mask[:,None], source=newsource)
-
-        non_finite_sum=(torch.isfinite(prev)==False).sum()
-        if(non_finite_sum>0):
-
-
-            print("NONZERO")
-            print((torch.isfinite(prev)==False).nonzero())
-
-
-            print("prev", prev[torch.isfinite(prev)==False])
-            print("feval ", f_eval[torch.isfinite(prev)==False])
-            print("f grad eval ", f_prime_eval[torch.isfinite(prev)==False])
-
-            raise Exception()
-
-        new_tolerance_mask=(torch.abs(update).sum(axis=1))>=newton_tolerance
-    
-        above_tolerance_mask=torch.masked_scatter(input=above_tolerance_mask, mask=above_tolerance_mask, source=new_tolerance_mask)
-
-        above_tol=above_tolerance_mask.sum()
-
-        if(verbose):
-            print("-- newton iter %d .. %d / %d dims completed" % (i, target_arg.shape[0]-above_tol, target_arg.shape[0]))
-        if(above_tol==0):
-            if(verbose):
-                print("------ done")
-            break
+        fn_result, f_prime_eval = joint_func(prev, *args)
 
-    if(target_arg.dtype==torch.float64):
+        f_eval = fn_result - target_arg
+        update = f_eval / f_prime_eval
 
-        target_prec=1e-7
-    else:
+        newsource = prev - update
+        prev = torch.where(active_row.unsqueeze(-1), newsource, prev)
 
-        target_prec=1e-4
+        still_active = torch.abs(update).sum(dim=1) >= newton_tolerance
+        active_row = active_row & still_active
 
-    num_non_converged=(torch.abs(f_eval)>target_prec).sum()
-    
-    if( num_non_converged>0):
-        print(num_non_converged, " items did not converge in Newton iterations")
-        print("feval (diff) ",f_eval[torch.abs(f_eval)>target_prec])
-    
     return prev
 
 def inverse_bisection_n_newton(func, 
diff --git a/jammy_flows/layers/spline_fns.py b/jammy_flows/layers/spline_fns.py
index 799d943..c670a30 100644
--- a/jammy_flows/layers/spline_fns.py
+++ b/jammy_flows/layers/spline_fns.py
@@ -85,12 +85,18 @@ def rational_quadratic_spline(inputs,
         heights = cumheights[..., 1:] - cumheights[..., :-1]
 
 
-        
+
         if inverse:
             bin_idx = searchsorted(cumheights, inputs)#[..., None]
         else:
             bin_idx = searchsorted(cumwidths, inputs)#[..., None]
-        
+
+        # Guard against NaN inputs (which yield -1 from searchsorted) or fp32
+        # boundary cases above/below the table.  Clamping keeps gather() safe;
+        # a NaN input still propagates through to a NaN output via the
+        # `inputs - input_cumheights` terms, without a device fault.
+        bin_idx = bin_idx.clamp(min=0, max=num_bins - 1)
+
         if(cumwidths.shape[0]==1 and bin_idx.shape[0]>1):
           repeats=[bin_idx.shape[0]]+(len(cumwidths.shape)-1)*[1]
           
@@ -130,7 +136,10 @@ def rational_quadratic_spline(inputs,
             c = - input_delta * (inputs - input_cumheights)
 
             discriminant = b.pow(2) - 4 * a * c
-            assert (discriminant >= 0).all()
+            # Discriminant is mathematically non-negative; catastrophic
+            # cancellation in fp32 can produce small negatives.  Clamp
+            # before sqrt to avoid NaNs.
+            discriminant = discriminant.clamp(min=0)
 
             root = (2 * c) / (-b - torch.sqrt(discriminant))
             outputs = root * input_bin_widths + input_cumwidths
@@ -141,7 +150,8 @@ def rational_quadratic_spline(inputs,
             derivative_numerator = input_delta.pow(2) * (input_derivatives_plus_one * root.pow(2)
                                                          + 2 * input_delta * theta_one_minus_theta
                                                          + input_derivatives * (1 - root).pow(2))
-            logabsdet = torch.log(derivative_numerator) - 2 * torch.log(denominator)
+            # Protect log from near-zero values at fp32 precision.
+            logabsdet = torch.log(derivative_numerator.clamp(min=1e-8)) - 2 * torch.log(denominator.clamp(min=1e-8))
 
             return outputs, -logabsdet
         else:
@@ -157,8 +167,8 @@ def rational_quadratic_spline(inputs,
             derivative_numerator = input_delta.pow(2) * (input_derivatives_plus_one * theta.pow(2)
                                                          + 2 * input_delta * theta_one_minus_theta
                                                          + input_derivatives * (1 - theta).pow(2))
-            logabsdet = torch.log(derivative_numerator) - 2 * torch.log(denominator)
-           
+            logabsdet = torch.log(derivative_numerator.clamp(min=1e-8)) - 2 * torch.log(denominator.clamp(min=1e-8))
+
             return outputs, logabsdet
 
 def rational_quadratic_spline_with_linear_extension(inputs,

From d3631f55b7aaad2b938209d9c65177f6a87ef454 Mon Sep 17 00:00:00 2001
From: JamesPearman230 <james@gridraven.com>
Date: Fri, 17 Apr 2026 17:02:33 +0300
Subject: [PATCH 04/11] Clean up comments

---
 jammy_flows/layers/bisection_n_newton.py | 15 +++++----------
 1 file changed, 5 insertions(+), 10 deletions(-)

diff --git a/jammy_flows/layers/bisection_n_newton.py b/jammy_flows/layers/bisection_n_newton.py
index 5617ac8..03aa2dd 100644
--- a/jammy_flows/layers/bisection_n_newton.py
+++ b/jammy_flows/layers/bisection_n_newton.py
@@ -56,15 +56,10 @@ def inverse_bisection_n_newton_joint_func_and_grad(func,
 
     prev = mid
 
-    # Newton iterations — compile-friendly variant.  The original code sliced
-    # `prev[above_tolerance_mask, :]` each iter to only operate on unconverged
-    # rows, but boolean-mask indexing produces a tensor with an unbacked symint
-    # size which breaks `torch.compile(fullgraph=True)`.  Instead, compute on
-    # the full batch every iteration and gate the update with `torch.where`
-    # so converged rows stay frozen.  Fixed `num_newton_iter` is done — no
-    # data-dependent early exit.  For our use case (D=1, num_newton_iter=20)
-    # the extra compute on converged rows is negligible and compile-trace
-    # fusion more than makes up for it.
+    # Changes to the Newton iterations for compile:
+    # 1. Keep the full batch each iteration, using torch.where to only update where
+    #    non-converged (used to slice the batch in 'prev[above_tolerance_mask, :]')
+    # 2. Removed the early exiting when converging so full number of iterations are completed
     active_row = torch.ones(target_arg.shape[0], dtype=torch.bool, device=target_arg.device)
     for i in range(num_newton_iter):
         fn_result, f_prime_eval = joint_func(prev, *args)
@@ -408,4 +403,4 @@ def inverse_bisection_n_newton_sphere_fast(combined_func,
             
             break
        
-    return prev
\ No newline at end of file
+    return prev

From 12e4734f7478860166961f6ac3bfa24db594b545 Mon Sep 17 00:00:00 2001
From: JamesPearman230 <james@gridraven.com>
Date: Thu, 25 Jun 2026 19:38:35 +0300
Subject: [PATCH 05/11] Safe denominator in
 rational_quadratic_spline_with_linear_extension

The forward (training/log_pdf) branch computes the RQ spline over the whole
tensor then overwrites the linear-tail region via torch.where. autograd
backprops through BOTH where-branches, so the dead RQ branch is still
evaluated at extrapolated theta for linear-tail rows, where denominator and
derivative_numerator can be <=0 -> inf/NaN -> NaN gradient (even though the
forward picks the finite linear branch).

clamp_min(1e-8) on denominator and derivative_numerator keeps the dead branch
finite. In-range rows are >> 1e-8 so forward output + log-det are unchanged.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
---
 jammy_flows/layers/spline_fns.py | 7 +++++++
 1 file changed, 7 insertions(+)

diff --git a/jammy_flows/layers/spline_fns.py b/jammy_flows/layers/spline_fns.py
index c670a30..ffab2a8 100644
--- a/jammy_flows/layers/spline_fns.py
+++ b/jammy_flows/layers/spline_fns.py
@@ -322,11 +322,18 @@ def rational_quadratic_spline_with_linear_extension(inputs,
                                          + input_derivatives * theta_one_minus_theta)
             denominator = input_delta + ((input_derivatives + input_derivatives_plus_one - 2 * input_delta)
                                          * theta_one_minus_theta)
+            # Safe denominator: for inputs in the linear tail (inputs<=left / >=right) the
+            # torch.where below overwrites outputs/logabsdet, but autograd still backprops
+            # through this RQ branch evaluated at extrapolated theta, where denominator and
+            # derivative_numerator can be <=0 -> inf/NaN in the dead branch -> NaN gradient.
+            # clamp_min keeps the dead branch finite; in-range rows are >> 1e-8 so unaffected.
+            denominator = denominator.clamp_min(1e-8)
             outputs = input_cumheights + numerator / denominator
 
             derivative_numerator = input_delta.pow(2) * (input_derivatives_plus_one * theta.pow(2)
                                                          + 2 * input_delta * theta_one_minus_theta
                                                          + input_derivatives * (1 - theta).pow(2))
+            derivative_numerator = derivative_numerator.clamp_min(1e-8)
             logabsdet = torch.log(derivative_numerator) - 2 * torch.log(denominator)
 
             ## fill in linear bits

From d5054b5209b01fab610a19d7c82ccadda833dd30 Mon Sep 17 00:00:00 2001
From: James Pearman <james@gridraven.com>
Date: Wed, 8 Jul 2026 13:48:09 +0300
Subject: [PATCH 06/11] fix bounds to stop nans in fp32

---
 jammy_flows/layers/spheres/sphere_base.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/jammy_flows/layers/spheres/sphere_base.py b/jammy_flows/layers/spheres/sphere_base.py
index 9f46f91..0ba59df 100644
--- a/jammy_flows/layers/spheres/sphere_base.py
+++ b/jammy_flows/layers/spheres/sphere_base.py
@@ -197,7 +197,7 @@ def eucl_to_spherical_embedding(self, x, log_det):
                     sq = torch.sum(x[:, ind:] ** 2, dim=1, keepdims=True)
                     denom = torch.sqrt(torch.clamp(sq, min=1e-8))
                     arg = x[:, ind:ind+1] / denom
-                    arg = torch.clamp(arg, -1.0, 1.0)
+                    arg = torch.clamp(arg, -1.0 + 1e-4, 1.0 - 1e-4)
                     new_angle = torch.acos(arg)
 
                     #mask_smaller=(x[:,ind+1:ind+2]<0).double()

From b8863d87a40930ccd52256346f526f1eb88cf3cc Mon Sep 17 00:00:00 2001
From: James Pearman <james@gridraven.com>
Date: Tue, 21 Jul 2026 10:59:48 +0300
Subject: [PATCH 07/11] bound discriminant to avoid nans in outputs

---
 jammy_flows/layers/spline_fns.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/jammy_flows/layers/spline_fns.py b/jammy_flows/layers/spline_fns.py
index ffab2a8..67f7a56 100644
--- a/jammy_flows/layers/spline_fns.py
+++ b/jammy_flows/layers/spline_fns.py
@@ -288,7 +288,7 @@ def rational_quadratic_spline_with_linear_extension(inputs,
                                                   - 2 * input_delta))
             c = - input_delta * (inputs - input_cumheights)
 
-            discriminant = b.pow(2) - 4 * a * c
+            discriminant = (b.pow(2) - 4 * a * c).clamp(min=0)
             #assert (discriminant >= 0).all(), (inputs[discriminant<0], input_cumwidths[discriminant<0], input_cumheights[discriminant<0], input_bin_widths[discriminant<0], input_heights[discriminant<0],bin_idx[discriminant<0],discriminant[discriminant<0], a[discriminant<0], b[discriminant<0], c[discriminant<0], a,b,c )
 
             root = (2 * c) / (-b - torch.sqrt(discriminant))

From fc67c4a2dc606ee19f2656ee478711aa17f22082 Mon Sep 17 00:00:00 2001
From: JamesPearman230 <james@gridraven.com>
Date: Wed, 22 Jul 2026 12:03:16 +0000
Subject: [PATCH 08/11] Re-apply fork optimizations onto the upstream sync

Restores this fork's torch.compile support and fp32 stability work on
top of the upstream merge, adapting it to upstream's new code:

- bisection_n_newton (inverse_bisection_n_newton_joint_func_and_grad):
  restore the compile-friendly full-batch Newton loop (torch.full_like,
  torch.where row masking, fixed iteration count, no data-dependent
  branching/prints). Upstream's new NaN/inf recovery is kept, folded in
  as an unconditional torch.where so fullgraph compilation still works:
  a non-finite Newton update falls back to the previous iterate.

- main/default.py: restore device-native sampling (torch.manual_seed +
  torch.randn instead of numpy round-trips), drop the removed debug
  blocks and the hot-path finiteness assertion in the sampling
  direction. Upstream's new only_last support in the forward/inverse
  layer loops is kept, folded into the fork's flattened loop style
  (python-level branching on static values, compile-safe).

- spline_fns: fix used_safety_margin bug in upstream's new
  return_safe_angle_within_2pi (bounds used the raw argument, which
  TypeErrors when None is passed) and widen its margin in fp32, where
  a sub-ULP margin at 2*pi made the upper clamp a silent no-op.

- sphere_base: make the s1 embedding acos clamp dtype-aware (1e-6 fp32
  / 1e-12 fp64). The previous fixed 1e-4 margin distorted the
  euclidean->angle map near 0/2pi enough to break float64 sample<->eval
  roundtrips (upstream test_selfconsistency, s1 'o' flows); the new
  margins keep NaN/inf protection in forward and backward passes while
  restoring roundtrip error to <2e-8.

Verified: upstream test suite (test_general, test_spheres,
test_manifold_embedding_consistency, test_entropy_and_marginal_entropy)
passes except one pre-existing marginal s2/'c' CNF tolerance case that
is bit-identical to pristine upstream under fixed inputs; fp32+fp64
forward/backward, boundary stress and _obtain_sample paths verified on
CPU+CUDA including torch.compile(fullgraph=True).

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
---
 jammy_flows/layers/bisection_n_newton.py  | 101 +++++-----------------
 jammy_flows/layers/spheres/sphere_base.py |   9 +-
 jammy_flows/layers/spline_fns.py          |  12 ++-
 jammy_flows/main/default.py               |  60 +++----------
 4 files changed, 55 insertions(+), 127 deletions(-)

diff --git a/jammy_flows/layers/bisection_n_newton.py b/jammy_flows/layers/bisection_n_newton.py
index a4f09a1..135431c 100644
--- a/jammy_flows/layers/bisection_n_newton.py
+++ b/jammy_flows/layers/bisection_n_newton.py
@@ -40,98 +40,45 @@ def inverse_bisection_n_newton_joint_func_and_grad(func,
     TODO: Make implicit at some point.
 
     """
-    new_upper = torch.tensor(max_boundary).type(target_arg.dtype).repeat(*target_arg.shape).to(target_arg.device)
-    new_lower = torch.tensor(min_boundary).type(target_arg.dtype).repeat(*target_arg.shape).to(target_arg.device)
- 
-    mid=0
+    new_upper = torch.full_like(target_arg, max_boundary)
+    new_lower = torch.full_like(target_arg, min_boundary)
+
+    mid = (new_upper + new_lower) / 2.
     for i in range(num_bisection_iter):
         mid = (new_upper + new_lower) / 2.
-        #print("mid: ", mid)
         inverse_mid = func(mid, *args)
 
-        #print("MID", mid)
-        
-        right_part = (inverse_mid < target_arg).type(target_arg.dtype)
+        right_part = (inverse_mid < target_arg).to(target_arg.dtype)
         left_part = 1. - right_part
 
-        correct_part = (close(inverse_mid, target_arg, rtol=1e-6, atol=0)).type(target_arg.dtype)
+        correct_part = close(inverse_mid, target_arg, rtol=1e-6, atol=0).to(target_arg.dtype)
 
         new_lower = (1. - correct_part) * (right_part * mid + left_part * new_lower) + correct_part * mid
         new_upper = (1. - correct_part) * (right_part * new_upper + left_part * mid) + correct_part * mid
-      
-        
-    prev=mid
-
-    #print("target arg", target_arg.shape)
 
+    prev = mid
 
-    above_tolerance_mask=torch.ones( target_arg.shape[0], dtype=torch.bool, device=target_arg.device)
-
-    ## check where we want to broadcast the masking, and wnhere not
-
-    broadcasting_bool_args=[True if (prev.shape[0]>1 and arg.shape[0]>1) else False for arg in args ]
-
+    # Changes to the Newton iterations for compile:
+    # 1. Keep the full batch each iteration, using torch.where to only update where
+    #    non-converged (used to slice the batch in 'prev[above_tolerance_mask, :]')
+    # 2. Removed the early exiting when converging so full number of iterations are completed
+    active_row = torch.ones(target_arg.shape[0], dtype=torch.bool, device=target_arg.device)
     for i in range(num_newton_iter):
-        
-        fn_result, f_prime_eval = joint_func(prev[above_tolerance_mask,:], *[a[above_tolerance_mask] if(broadcasting_bool_args[arg_index] == True) else a for arg_index, a in enumerate(args)])
-        
-        f_eval=fn_result-target_arg[above_tolerance_mask,:]
-
-        update=(f_eval/f_prime_eval)
-
-        newsource=prev[above_tolerance_mask,:]-update
-       
-        ##### correction if we get nans/infs
-        non_fin_mask_new=~torch.isfinite(newsource)
-        non_finite_sum=(non_fin_mask_new==True).sum()
-        if(non_finite_sum>0):
-            print("---- non finite in jammy flows sampling.. try to FIX IT .....................")
-            # stop iterations and replace with previous
-            newsource=torch.where(non_fin_mask_new, prev[above_tolerance_mask,:], newsource)
-        ##################
-
-        prev=torch.masked_scatter(input=prev, mask=above_tolerance_mask[:,None], source=newsource)
-
-        non_finite_sum=(torch.isfinite(prev)==False).sum()
-        if(non_finite_sum>0):
-
-
-            print("NONZERO")
-            print((torch.isfinite(prev)==False).nonzero())
+        fn_result, f_prime_eval = joint_func(prev, *args)
 
+        f_eval = fn_result - target_arg
+        update = f_eval / f_prime_eval
 
-            print("prev", prev[torch.isfinite(prev)==False])
-            print("feval ", f_eval[torch.isfinite(prev)==False])
-            print("f grad eval ", f_prime_eval[torch.isfinite(prev)==False])
+        newsource = prev - update
+        # Upstream NaN/inf correction (folded in unconditionally to stay
+        # compile-friendly): a non-finite Newton update falls back to the
+        # previous iterate instead of poisoning the result.
+        newsource = torch.where(torch.isfinite(newsource), newsource, prev)
+        prev = torch.where(active_row.unsqueeze(-1), newsource, prev)
 
-            raise Exception()
+        still_active = torch.abs(update).sum(dim=1) >= newton_tolerance
+        active_row = active_row & still_active
 
-        new_tolerance_mask=(torch.abs(update).sum(axis=1))>=newton_tolerance
-    
-        above_tolerance_mask=torch.masked_scatter(input=above_tolerance_mask, mask=above_tolerance_mask, source=new_tolerance_mask)
-
-        above_tol=above_tolerance_mask.sum()
-
-        if(verbose):
-            print("-- newton iter %d .. %d / %d dims completed" % (i, target_arg.shape[0]-above_tol, target_arg.shape[0]))
-        if(above_tol==0):
-            if(verbose):
-                print("------ done")
-            break
-
-    if(target_arg.dtype==torch.float64):
-
-        target_prec=1e-7
-    else:
-
-        target_prec=1e-4
-
-    num_non_converged=(torch.abs(f_eval)>target_prec).sum()
-    
-    if( num_non_converged>0):
-        print(num_non_converged, " items did not converge in Newton iterations")
-        print("feval (diff) ",f_eval[torch.abs(f_eval)>target_prec])
-    
     return prev
 
 def inverse_bisection_n_newton(func, 
@@ -462,4 +409,4 @@ def inverse_bisection_n_newton_sphere_fast(combined_func,
             
             break
        
-    return prev
\ No newline at end of file
+    return prev
diff --git a/jammy_flows/layers/spheres/sphere_base.py b/jammy_flows/layers/spheres/sphere_base.py
index 9eebd9b..db2d4d9 100644
--- a/jammy_flows/layers/spheres/sphere_base.py
+++ b/jammy_flows/layers/spheres/sphere_base.py
@@ -258,10 +258,17 @@ def eucl_to_spherical_embedding(self, x, log_det):
 
                     # last one is 0 to 2pi
                     # new_angle=torch.acos(x[:,ind:ind+1]/torch.sum(x[:,ind:]**2, dim=1, keepdims=True).sqrt())
+                    # Keep acos away from +-1 so neither the value nor the
+                    # backward pass (d/dx acos = -1/sqrt(1-x^2)) produces
+                    # NaN/inf. Margin is dtype-aware: wide enough to survive
+                    # fp32 rounding, tight enough not to distort the inverse
+                    # mapping in fp64 (a fixed 1e-4 margin breaks
+                    # sample<->eval roundtrip consistency).
+                    acos_margin = 1e-6 if x.dtype == torch.float32 else 1e-12
                     sq = torch.sum(x[:, ind:] ** 2, dim=1, keepdims=True)
                     denom = torch.sqrt(torch.clamp(sq, min=1e-8))
                     arg = x[:, ind:ind+1] / denom
-                    arg = torch.clamp(arg, -1.0 + 1e-4, 1.0 - 1e-4)
+                    arg = torch.clamp(arg, -1.0 + acos_margin, 1.0 - acos_margin)
                     new_angle = torch.acos(arg)
 
                     #mask_smaller=(x[:,ind+1:ind+2]<0).double()
diff --git a/jammy_flows/layers/spline_fns.py b/jammy_flows/layers/spline_fns.py
index d759be2..fff22bb 100644
--- a/jammy_flows/layers/spline_fns.py
+++ b/jammy_flows/layers/spline_fns.py
@@ -31,9 +31,15 @@ def return_safe_angle_within_2pi(x, safety_margin=1e-7):
             used_safety_margin=1e-7
         elif(x.dtype==torch.float64):
             used_safety_margin=1e-10
-    
-    upper_bound=2*numpy.pi-safety_margin
-    lower_bound=0.0+safety_margin
+
+    # In float32 the spacing between representable values at 2*pi is ~4.8e-7,
+    # so a smaller margin makes the upper clamp a no-op (2*pi - margin rounds
+    # back to 2*pi). Widen it so the clamp stays effective in fp32.
+    if(x.dtype==torch.float32 and used_safety_margin<1e-6):
+        used_safety_margin=1e-6
+
+    upper_bound=2*numpy.pi-used_safety_margin
+    lower_bound=0.0+used_safety_margin
 
     small_mask=x<lower_bound
     large_mask=x>upper_bound
diff --git a/jammy_flows/main/default.py b/jammy_flows/main/default.py
index 91c10b9..4a689e0 100644
--- a/jammy_flows/main/default.py
+++ b/jammy_flows/main/default.py
@@ -918,8 +918,6 @@ def all_layer_inverse(self,
         extra_conditional_input=[]
         base_targets=[]
 
-        individual_logps=dict()
-
         extra_params = None
 
         if(amortization_parameters is not None):
@@ -1011,40 +1009,18 @@ def all_layer_inverse(self,
                             - layer.total_param_num : -extra_param_counter,
                         ]
 
-                
-                if(l==(len(pdf_layers)-1)):
+                if(only_last and l==(len(pdf_layers)-1)):
                     # force embedding or intrinsic coordinates in the layer that defines the target dimension
-                    
-                    if(only_last):
-                        if(self.pdf_defs_list[pdf_index][0]=="s"):
-                            this_target, log_det = layer.inv_flow_mapping([this_target, log_det], extra_inputs=this_extra_params, fix_euclidean_to_sphere_first=True)
-                        else:
-                            this_target, log_det = layer.inv_flow_mapping([this_target, log_det], extra_inputs=this_extra_params)
-                        break
+                    if(self.pdf_defs_list[pdf_index][0]=="s"):
+                        this_target, log_det = layer.inv_flow_mapping([this_target, log_det], extra_inputs=this_extra_params, fix_euclidean_to_sphere_first=True)
                     else:
                         this_target, log_det = layer.inv_flow_mapping([this_target, log_det], extra_inputs=this_extra_params)
-
+                    break
                 else:
-
                     this_target, log_det = layer.inv_flow_mapping([this_target, log_det], extra_inputs=this_extra_params)
-                
-                extra_param_counter += layer.total_param_num
-
-            if(False):
-                ## stems from debugging purposes, not used currently
-                ind_base_eval=this_logp = torch.distributions.MultivariateNormal(
-                    torch.zeros_like(this_target).to(x),
-                    covariance_matrix=torch.eye(this_target.shape[1]).type_as(x).to(x),
-                ).log_prob(this_target)
-
-                ind_logdet=log_det
-                
 
-                individual_logps["%.2d_%s" % (pdf_index, this_pdf_type)]=ind_base_eval+ind_logdet
-                individual_logps["%.2d_%s_logdet" % (pdf_index, this_pdf_type)]=ind_logdet
-                individual_logps["%.2d_%s_base" % (pdf_index, this_pdf_type)]=ind_base_eval
+                extra_param_counter += layer.total_param_num
 
-            
             base_targets.append(this_target)
 
             prev_target=x[:,self.target_dim_indices[pdf_index][0]:self.target_dim_indices[pdf_index][1]]
@@ -1203,12 +1179,11 @@ def obtain_flow_param_structure(self,
         else:
 
             if(seed is not None):
-                numpy.random.seed(seed)
-
-            std_normal = numpy.random.normal(size=(used_sample_size, self.total_base_dim))
+                torch.manual_seed(seed)
 
-            std_normal_samples = (
-                torch.from_numpy(std_normal).type(data_type).to(used_device)
+            std_normal_samples = torch.randn(
+                used_sample_size, self.total_base_dim,
+                dtype=data_type, device=used_device,
             )
             log_gauss_evals = torch.distributions.MultivariateNormal(
                 torch.zeros(self.total_base_dim).type(data_type).to(used_device),
@@ -1486,14 +1461,11 @@ def all_layer_forward(self,
                 if extra_params is not None:
                     
                     this_extra_params = extra_params[:, extra_param_counter : extra_param_counter + layer.total_param_num]
-
                 if(only_last):
                     if(l<(len(pdf_layers)-1) ):
                         extra_param_counter += layer.total_param_num
                         continue
 
-                if(only_last):
-                 
                     if(self.pdf_defs_list[pdf_index][0]=="s"):
                         this_target, log_det = layer.flow_mapping([this_target, log_det], extra_inputs=this_extra_params, fix_euclidean_to_sphere_first=True)
                     elif(self.pdf_defs_list[pdf_index][0]=="e"):
@@ -1502,7 +1474,7 @@ def all_layer_forward(self,
                         raise Exception("Flow type ", self.pdf_defs_list[pdf_index][0], " does not supported *only_last*!")
                 else:
                     this_target, log_det = layer.flow_mapping([this_target, log_det], extra_inputs=this_extra_params)
-              
+
                 extra_param_counter += layer.total_param_num
 
             new_targets.append(this_target)
@@ -1513,9 +1485,6 @@ def all_layer_forward(self,
 
             extra_conditional_input.append(prev_target)
 
-        if (torch.isfinite(x) == 0).sum() > 0:
-            raise Exception("nonfinite samples generated .. this should never happen!")
-        
         x=torch.cat(new_targets, dim=1)
 
         ## transform to desired output space 
@@ -1659,12 +1628,11 @@ def _obtain_sample(self,
         else:
 
             if(seed is not None):
-                numpy.random.seed(seed)
-
-            std_normal = numpy.random.normal(size=(used_sample_size, self.total_base_dim))
+                torch.manual_seed(seed)
 
-            std_normal_samples = (
-                torch.from_numpy(std_normal).type(data_type).to(used_device)
+            std_normal_samples = torch.randn(
+                used_sample_size, self.total_base_dim,
+                dtype=data_type, device=used_device,
             )
 
             log_gauss_evals=torch.distributions.Normal(0.0,1.0).log_prob(std_normal_samples).sum(dim=-1)

From 2288ade0a3f885e45a26b59410d91b6dc14b4c54 Mon Sep 17 00:00:00 2001
From: JamesPearman230 <james@gridraven.com>
Date: Mon, 27 Jul 2026 15:06:34 +0300
Subject: [PATCH 09/11] Update sphere_base.py

---
 jammy_flows/layers/spheres/sphere_base.py | 10 +++-------
 1 file changed, 3 insertions(+), 7 deletions(-)

diff --git a/jammy_flows/layers/spheres/sphere_base.py b/jammy_flows/layers/spheres/sphere_base.py
index db2d4d9..72d1a3b 100644
--- a/jammy_flows/layers/spheres/sphere_base.py
+++ b/jammy_flows/layers/spheres/sphere_base.py
@@ -258,12 +258,8 @@ def eucl_to_spherical_embedding(self, x, log_det):
 
                     # last one is 0 to 2pi
                     # new_angle=torch.acos(x[:,ind:ind+1]/torch.sum(x[:,ind:]**2, dim=1, keepdims=True).sqrt())
-                    # Keep acos away from +-1 so neither the value nor the
-                    # backward pass (d/dx acos = -1/sqrt(1-x^2)) produces
-                    # NaN/inf. Margin is dtype-aware: wide enough to survive
-                    # fp32 rounding, tight enough not to distort the inverse
-                    # mapping in fp64 (a fixed 1e-4 margin breaks
-                    # sample<->eval roundtrip consistency).
+                    # keep acos away from |1| so no nans/inf in forward or backward
+                    # functionality is dependent on dtype...
                     acos_margin = 1e-6 if x.dtype == torch.float32 else 1e-12
                     sq = torch.sum(x[:, ind:] ** 2, dim=1, keepdims=True)
                     denom = torch.sqrt(torch.clamp(sq, min=1e-8))
@@ -887,4 +883,4 @@ def _flow_mapping(self, inputs, extra_inputs=None, sf_extra=None):
 
 
 
-    
\ No newline at end of file
+    

From e441a5dda26a64ce2e3bb113c156dfa8ab8356a5 Mon Sep 17 00:00:00 2001
From: James Pearman <james@gridraven.com>
Date: Tue, 28 Jul 2026 15:35:52 +0300
Subject: [PATCH 10/11] numerical clamp changes to address regression tests

---
 jammy_flows/layers/spheres/sphere_base.py | 12 ++++--------
 jammy_flows/layers/spline_fns.py          | 18 +++++++-----------
 jammy_flows/main/default.py               | 14 ++++++++++----
 3 files changed, 21 insertions(+), 23 deletions(-)

diff --git a/jammy_flows/layers/spheres/sphere_base.py b/jammy_flows/layers/spheres/sphere_base.py
index 72d1a3b..a0f6262 100644
--- a/jammy_flows/layers/spheres/sphere_base.py
+++ b/jammy_flows/layers/spheres/sphere_base.py
@@ -258,14 +258,10 @@ def eucl_to_spherical_embedding(self, x, log_det):
 
                     # last one is 0 to 2pi
                     # new_angle=torch.acos(x[:,ind:ind+1]/torch.sum(x[:,ind:]**2, dim=1, keepdims=True).sqrt())
-                    # keep acos away from |1| so no nans/inf in forward or backward
-                    # functionality is dependent on dtype...
-                    acos_margin = 1e-6 if x.dtype == torch.float32 else 1e-12
-                    sq = torch.sum(x[:, ind:] ** 2, dim=1, keepdims=True)
-                    denom = torch.sqrt(torch.clamp(sq, min=1e-8))
-                    arg = x[:, ind:ind+1] / denom
-                    arg = torch.clamp(arg, -1.0 + acos_margin, 1.0 - acos_margin)
-                    new_angle = torch.acos(arg)
+                    # atan2 replaces acos (exact at +-1 w finite grads
+                    rest_sq = torch.sum(x[:, ind+1:] ** 2, dim=1, keepdims=True)
+                    rest_norm = torch.sqrt(torch.clamp(rest_sq, min=torch.finfo(x.dtype).tiny))
+                    new_angle = torch.atan2(rest_norm, x[:, ind:ind+1])
 
                     #mask_smaller=(x[:,ind+1:ind+2]<0).double()
                     new_angle=torch.where(x[:,ind+1:ind+2]<0, 2*numpy.pi-new_angle, new_angle)
diff --git a/jammy_flows/layers/spline_fns.py b/jammy_flows/layers/spline_fns.py
index fff22bb..1f2e5a0 100644
--- a/jammy_flows/layers/spline_fns.py
+++ b/jammy_flows/layers/spline_fns.py
@@ -177,8 +177,9 @@ def rational_quadratic_spline(inputs,
             derivative_numerator = input_delta.pow(2) * (input_derivatives_plus_one * root.pow(2)
                                                          + 2 * input_delta * theta_one_minus_theta
                                                          + input_derivatives * (1 - root).pow(2))
-            # Protect log from near-zero values at fp32 precision.
-            logabsdet = torch.log(derivative_numerator.clamp(min=1e-8)) - 2 * torch.log(denominator.clamp(min=1e-8))
+            # Floor at min dtype normal
+            log_floor = torch.finfo(derivative_numerator.dtype).tiny
+            logabsdet = torch.log(derivative_numerator.clamp(min=log_floor)) - 2 * torch.log(denominator.clamp(min=log_floor))
 
             return outputs, -logabsdet
         else:
@@ -194,7 +195,8 @@ def rational_quadratic_spline(inputs,
             derivative_numerator = input_delta.pow(2) * (input_derivatives_plus_one * theta.pow(2)
                                                          + 2 * input_delta * theta_one_minus_theta
                                                          + input_derivatives * (1 - theta).pow(2))
-            logabsdet = torch.log(derivative_numerator.clamp(min=1e-8)) - 2 * torch.log(denominator.clamp(min=1e-8))
+            log_floor = torch.finfo(derivative_numerator.dtype).tiny
+            logabsdet = torch.log(derivative_numerator.clamp(min=log_floor)) - 2 * torch.log(denominator.clamp(min=log_floor))
 
             return outputs, logabsdet
 
@@ -342,25 +344,19 @@ def rational_quadratic_spline_with_linear_extension(inputs,
 
             return outputs, final_logabsdet
         else:
-            theta = (inputs - input_cumwidths) / input_bin_widths
+            # clamp to avoid inf/NaN grads, avoids needing clamp for denom later
+            theta = ((inputs - input_cumwidths) / input_bin_widths).clamp(0.0, 1.0)
             theta_one_minus_theta = theta * (1 - theta)
 
             numerator = input_heights * (input_delta * theta.pow(2)
                                          + input_derivatives * theta_one_minus_theta)
             denominator = input_delta + ((input_derivatives + input_derivatives_plus_one - 2 * input_delta)
                                          * theta_one_minus_theta)
-            # Safe denominator: for inputs in the linear tail (inputs<=left / >=right) the
-            # torch.where below overwrites outputs/logabsdet, but autograd still backprops
-            # through this RQ branch evaluated at extrapolated theta, where denominator and
-            # derivative_numerator can be <=0 -> inf/NaN in the dead branch -> NaN gradient.
-            # clamp_min keeps the dead branch finite; in-range rows are >> 1e-8 so unaffected.
-            denominator = denominator.clamp_min(1e-8)
             outputs = input_cumheights + numerator / denominator
 
             derivative_numerator = input_delta.pow(2) * (input_derivatives_plus_one * theta.pow(2)
                                                          + 2 * input_delta * theta_one_minus_theta
                                                          + input_derivatives * (1 - theta).pow(2))
-            derivative_numerator = derivative_numerator.clamp_min(1e-8)
             logabsdet = torch.log(derivative_numerator) - 2 * torch.log(denominator)
 
             ## fill in linear bits
diff --git a/jammy_flows/main/default.py b/jammy_flows/main/default.py
index 4a689e0..ba6b439 100644
--- a/jammy_flows/main/default.py
+++ b/jammy_flows/main/default.py
@@ -1178,12 +1178,15 @@ def obtain_flow_param_structure(self,
 
         else:
 
+            # local generator: seeding must not perturb the caller's global RNG stream
+            generator = None
             if(seed is not None):
-                torch.manual_seed(seed)
+                generator = torch.Generator(device=used_device)
+                generator.manual_seed(seed)
 
             std_normal_samples = torch.randn(
                 used_sample_size, self.total_base_dim,
-                dtype=data_type, device=used_device,
+                dtype=data_type, device=used_device, generator=generator,
             )
             log_gauss_evals = torch.distributions.MultivariateNormal(
                 torch.zeros(self.total_base_dim).type(data_type).to(used_device),
@@ -1627,12 +1630,15 @@ def _obtain_sample(self,
 
         else:
 
+            # local generator: seeding must not perturb the caller's global RNG stream
+            generator = None
             if(seed is not None):
-                torch.manual_seed(seed)
+                generator = torch.Generator(device=used_device)
+                generator.manual_seed(seed)
 
             std_normal_samples = torch.randn(
                 used_sample_size, self.total_base_dim,
-                dtype=data_type, device=used_device,
+                dtype=data_type, device=used_device, generator=generator,
             )
 
             log_gauss_evals=torch.distributions.Normal(0.0,1.0).log_prob(std_normal_samples).sum(dim=-1)

From 9b5afdb2c959bbc171d58e2eb5740b0c06ca96dc Mon Sep 17 00:00:00 2001
From: James Pearman <james@gridraven.com>
Date: Wed, 29 Jul 2026 14:38:00 +0300
Subject: [PATCH 11/11] make generator a single op node to avoid compilation
 breaks

---
 jammy_flows/main/default.py | 21 +++++----------------
 jammy_flows/rng_fns.py      | 30 ++++++++++++++++++++++++++++++
 2 files changed, 35 insertions(+), 16 deletions(-)
 create mode 100644 jammy_flows/rng_fns.py

diff --git a/jammy_flows/main/default.py b/jammy_flows/main/default.py
index ba6b439..7df8186 100644
--- a/jammy_flows/main/default.py
+++ b/jammy_flows/main/default.py
@@ -17,6 +17,7 @@
 from ..flow_options import check_flow_option, obtain_default_options, obtain_overall_flow_info
 from ..extra_functions import list_from_str, NONLINEARITIES, recheck_sampling, find_init_pars_of_chained_blocks
 from ..amortizable_mlp import AmortizableMLP
+from ..rng_fns import draw_standard_normal
 from ..helper_fns import contours, grid_functions
 from ..helper_fns.coverage import calculate_approximate_coverage
 from ..helper_fns.plotting.spherical import get_multiresolution_evals,plot_multiresolution_healpy
@@ -1178,15 +1179,9 @@ def obtain_flow_param_structure(self,
 
         else:
 
-            # local generator: seeding must not perturb the caller's global RNG stream
-            generator = None
-            if(seed is not None):
-                generator = torch.Generator(device=used_device)
-                generator.manual_seed(seed)
-
-            std_normal_samples = torch.randn(
+            std_normal_samples = draw_standard_normal(
                 used_sample_size, self.total_base_dim,
-                dtype=data_type, device=used_device, generator=generator,
+                data_type, used_device, seed=seed,
             )
             log_gauss_evals = torch.distributions.MultivariateNormal(
                 torch.zeros(self.total_base_dim).type(data_type).to(used_device),
@@ -1630,15 +1625,9 @@ def _obtain_sample(self,
 
         else:
 
-            # local generator: seeding must not perturb the caller's global RNG stream
-            generator = None
-            if(seed is not None):
-                generator = torch.Generator(device=used_device)
-                generator.manual_seed(seed)
-
-            std_normal_samples = torch.randn(
+            std_normal_samples = draw_standard_normal(
                 used_sample_size, self.total_base_dim,
-                dtype=data_type, device=used_device, generator=generator,
+                data_type, used_device, seed=seed,
             )
 
             log_gauss_evals=torch.distributions.Normal(0.0,1.0).log_prob(std_normal_samples).sum(dim=-1)
diff --git a/jammy_flows/rng_fns.py b/jammy_flows/rng_fns.py
new file mode 100644
index 0000000..6990e03
--- /dev/null
+++ b/jammy_flows/rng_fns.py
@@ -0,0 +1,30 @@
+import torch
+# custom op so that fullgraph compilation traces this as single node
+
+
+def _seeded_standard_normal(num_samples: int,
+                            dim: int,
+                            seed: int,
+                            dtype: torch.dtype,
+                            device: torch.device) -> torch.Tensor:
+
+    generator = torch.Generator(device=device)
+    generator.manual_seed(seed)
+
+    return torch.randn(num_samples, dim, dtype=dtype, device=device, generator=generator)
+
+
+# custom_op requires torch>=2.4
+if(hasattr(torch.library, "custom_op")):
+    _seeded_standard_normal = torch.library.custom_op(
+        "jammy_flows::seeded_standard_normal", mutates_args=())(_seeded_standard_normal)
+
+    _seeded_standard_normal.register_fake(
+        lambda num_samples, dim, seed, dtype, device: torch.empty(num_samples, dim, dtype=dtype, device=device))
+
+
+def draw_standard_normal(num_samples, dim, dtype, device, seed=None):
+    if(seed is None):
+        return torch.randn(num_samples, dim, dtype=dtype, device=device)
+
+    return _seeded_standard_normal(num_samples, dim, int(seed), dtype, device)