Add differentiable Procrustes registration#1818
Draft
loliverhennigh wants to merge 1 commit into
Draft
Conversation
loliverhennigh
force-pushed
the
codex/rigid-registration
branch
2 times, most recently
from
July 9, 2026 21:22
1466a0a to
53240a9
Compare
Signed-off-by: Oliver Hennigh <loliverhennigh101@gmail.com>
loliverhennigh
force-pushed
the
codex/rigid-registration
branch
from
July 9, 2026 21:46
53240a9 to
0759ed7
Compare
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
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.
Summary
physicsnemo.nn.functional.procrustesundergeometry.deformfor orientation-preserving rigid and similarity registration of point sets with known correspondences.(N, D)and aligned batched(B, N, D)inputs forDin{1, 2, 3}, with optional nonnegative isotropic scale estimation.torch.library.custom_op, fake registration, and registered autograd.opcheck, compile, CUDA Graph, dispatch, benchmark-contract, validation, and numerical edge-case tests, plus API docs, changelog coverage, and a functional visualization.Why
This is the closed-form Procrustes/Kabsch portion of rigid-registration roadmap item NVIDIA/physicsnemo-roadmap#2452. It supplies a differentiable alignment primitive for corresponding point sets while following the deformation package and backend conventions established by #1786.
PR #1786 has now merged, so this branch is rebased directly onto its merge commit on
mainand the diff contains only the Procrustes work.Scope
This PR intentionally implements Procrustes only. ICP remains a separate follow-up because
soft=Truedoes not yet define the correspondence relaxation, temperature, iteration and stopping policy, outlier handling, scale policy, gradient semantics, or transform return type.The returned transform follows the row-vector convention:
Implementation notes
Validation
pytest -q test/nn/functional/geometry/test_point_procrustes.py: 93 passed, 5 skipped (CUDA-only)pre-commit run --from-ref origin/main --to-ref HEAD: all hooks passedopcheck, fake-tensor/AOT tracing, fullgraph AOT-eager forward/backward, finite-difference gradcheck, and unsupported second-derivative behavior passedCUDA hardware was unavailable locally, so no GPU performance claim is made; CUDA dispatch and graph-capture coverage remains in the test suite for CI.