Feat/fix generator compiling - #11
Closed
JamesPearman230 wants to merge 17 commits into
Closed
Conversation
…ding small changes to avoid numerical instability
…ding small changes to avoid numerical instability
Torch compile support
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>
Brings in upstream's 26 new commits: smooth rational-quadratic splines and a reworked 'o' circular-spline layer, quaternion/xyz rotation modes, safer bisection, only_last evaluation support, zlp-kent utilities, and removal of the outdated 'n'/'p' flows. The two conflicted files (bisection_n_newton.py, main/default.py) are resolved wholesale in UPSTREAM's favor in this commit; the follow-up commit re-applies this fork's torch.compile and fp32-stability adaptations on top, so upstream's changes and the fork adjustments stay separately reviewable. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
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>
Upstream sync
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
No description provided.