Feat/fix generator compiling

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
+    return prev

jammy_flows/layers/spheres/sphere_base.py

@@ -257,7 +257,12 @@ 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())
+                    # 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)
 
@@ -874,4 +879,4 @@ def _flow_mapping(self, inputs, extra_inputs=None, sf_extra=None):
 
 
 
-
+

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
@@ -30,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
@@ -54,10 +61,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:
@@ -109,12 +112,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]
 
@@ -154,7 +163,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
@@ -165,7 +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))
-            logabsdet = torch.log(derivative_numerator) - 2 * torch.log(denominator)
+            # 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:
@@ -181,8 +195,9 @@ 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)
-
+            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
 
 def rational_quadratic_spline_with_linear_extension(inputs,
@@ -302,7 +317,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))
@@ -329,7 +344,8 @@ 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)

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
@@ -918,8 +919,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 +1010,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]]
@@ -1202,13 +1179,9 @@ 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))
-
-            std_normal_samples = (
-                torch.from_numpy(std_normal).type(data_type).to(used_device)
+            std_normal_samples = draw_standard_normal(
+                used_sample_size, self.total_base_dim,
+                data_type, used_device, seed=seed,
             )
             log_gauss_evals = torch.distributions.MultivariateNormal(
                 torch.zeros(self.total_base_dim).type(data_type).to(used_device),
@@ -1486,14 +1459,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 +1472,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 +1483,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 
@@ -1658,13 +1625,9 @@ 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))
-
-            std_normal_samples = (
-                torch.from_numpy(std_normal).type(data_type).to(used_device)
+            std_normal_samples = draw_standard_normal(
+                used_sample_size, self.total_base_dim,
+                data_type, used_device, seed=seed,
             )
 
             log_gauss_evals=torch.distributions.Normal(0.0,1.0).log_prob(std_normal_samples).sum(dim=-1)