Mesh: normalize geometry robustly across dtypes and scales - #1883
Open
peterdsharpe wants to merge 2 commits into
Open
Mesh: normalize geometry robustly across dtypes and scales#1883peterdsharpe wants to merge 2 commits into
peterdsharpe wants to merge 2 commits 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 |
Signed-off-by: Peter Sharpe <peterdsharpe@gmail.com>
Collaborator
Author
|
/ok to test 56d6038 |
peterdsharpe
marked this pull request as ready for review
September 9, 2026 14:05
Contributor
  The PR appears safe to merge after addressing the non-blocking autocast dtype-preservation gap in Findings
|
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
| is_zero = scale == 0 | ||
| scaled = vectors / scale.masked_fill(is_zero, 1) | ||
| norm = scaled.norm(dim=dim, keepdim=True) | ||
| return scaled / norm.masked_fill(is_zero, 1) |
Contributor
There was a problem hiding this comment.
Under CUDA autocast, scaled.norm(...) can produce an fp32 tensor for bf16 or float16 input, causing this final division to return fp32. This conflicts with safe_normalize's documented promise to preserve the input dtype. compute_cell_normals explicitly casts its result back for this reason, but the point, partition, and transformed-normal callers do not, so they can silently pass fp32 normals to downstream consumers.
Suggested change
| return scaled / norm.masked_fill(is_zero, 1) | |
| return (scaled / norm.masked_fill(is_zero, 1)).to(vectors.dtype) |
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 original isolated benchmark reported a roughly 2–4x kernel cost for the extra scaling pass; that benchmark was not rerun for this revision. Normals remain cached per mesh. The guarantee concerns finite-input forward normalization; extreme intermediate derivatives can still exceed a dtype’s range.
Refreshed against current main while preserving
with_points/with_datacache handling and the synchronization-freesolve_ex(..., check_errors=False)path. Both cached-normal transformations use the new normalization helper.Verification
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.