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 8c11900..a0f6262 100644 --- a/jammy_flows/layers/spheres/sphere_base.py +++ b/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): - \ No newline at end of file + diff --git a/jammy_flows/layers/spline_fns.py b/jammy_flows/layers/spline_fns.py index 22d615d..1f2e5a0 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 @@ -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=xupper_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) diff --git a/jammy_flows/main/default.py b/jammy_flows/main/default.py index 91c10b9..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 @@ -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) 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)