Mesh: normalize geometry robustly across dtypes and scales - #1883
Draft
peterdsharpe wants to merge 1 commit into
Draft
Mesh: normalize geometry robustly across dtypes and scales#1883peterdsharpe wants to merge 1 commit into
peterdsharpe wants to merge 1 commit into
Conversation
The hardcoded eps=1e-12 clamp in torch.nn.functional.normalize is not safe across the dtypes and length scales mesh geometry uses. In float16 the floor itself rounds to zero, so degenerate cells divide 0/0 and return NaN. Forming the norm squares each component and overflows to inf for large cells, so a well-conditioned normal silently becomes zero. And because the floor is absolute rather than dtype-relative, a genuine norm below it is substituted, so float32 and float64 alike return non-unit normals once feature size falls below roughly 1e-6. Add safe_normalize, which divides each vector by its own largest absolute component before taking the norm and preserves exactly-zero vectors, and route cell, point, cluster, and transformed normals through it. The rescaled norm is bounded in [1, sqrt(n)], so it can neither overflow nor underflow and no epsilon is needed. Also replaces a boolean-mask centroid assignment in partition_cells with torch.where, which is bit-identical and drops a nonzero device sync.
Contributor
CODEOWNERS review mapCurrent for commit ⏳ @peterdsharpe — 9 file(s)
No CODEOWNER
Comment |
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
1 participant
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.
PhysicsNeMo Pull Request
Description
torch.nn.functional.normalize(eps=1e-12)is not scale- or dtype-safe for mesh normals: the floor rounds to zero infloat16and turns zero vectors into NaN, norm formation can overflow and turn valid normals into zero, and legitimate smallfloat32orfloat64vectors become non-unit normals.This PR:
safe_normalize, which first scales each vector by its largest absolute component so its norm cannot overflow or underflow.torch.wherelogic to avoid a device synchronization.The extra scaling pass makes the normalization kernel roughly 2-4x slower in isolation, but normals are cached per mesh and remain sub-millisecond at multi-million-cell scale; the tradeoff removes dependence on coordinate scale and dtype.
Verification
pytest test/mesh: 2370 passed, 701 skipped.F.normalizeimplementation and a zero-mask-only implementation that omits rescaling.float16,bfloat16,float32, andfloat64, including zero, underflowing, and overflowing vectors.ruff check,ruff format --check, and doctests pass.Checklist
Dependencies
None.
Review Process
All PRs are reviewed by the PhysicsNeMo team before merging.
Depending on which files are changed, GitHub may automatically assign a maintainer for review.
We are also testing AI-based code review tools (e.g., Greptile), which may add automated comments with a confidence score. This score reflects the AI's assessment of merge readiness and is not a qualitative judgment of the work or an indication that the PR will be accepted or rejected.
AI-generated feedback should be reviewed critically for usefulness. You are not required to respond to every AI comment, but they are intended to help both authors and reviewers. Please react to Greptile comments with 👍 or 👎 to provide feedback on their accuracy.