[poketopa]

procrustes_naive flattens the batch dimensions of M before the SVD and restores them for R. With return_singular_values=True, it returns the singular values DS with the flattened shape. procrustes restores the batch dimensions of both outputs, so the two functions disagree whenever M does not have exactly one batch dimension:

• M.shape == (3, 3): procrustes returns DS of shape (3,), procrustes_naive returns (1, 3).
• M.shape == (2, 5, 3, 3): procrustes returns (2, 5, 3), procrustes_naive returns (10, 3).

special_procrustes_naive calls procrustes_naive and has the same issue.

Restore the batch dimensions of DS with roma.internal.unflatten_batch_dims, as procrustes does. The new test runs procrustes_naive on inputs with zero, one and two batch dimensions, with and without force_rotation, checks the shapes of R and DS, and compares both outputs with procrustes.

Validation on current master 9b168674896f2b586cb8addb4fbc9ed69f1e7df4, CPU:

• The new test fails on unchanged master (torch.Size([1, 3]) != (3,)) and passes with the fix.
• python -m unittest: all 72 tests pass with CPython 3.12.12, PyTorch 2.14.1 and NumPy 2.5.4. The suite also passes with the two older CI configurations, CPython 3.12.12 / PyTorch 2.3.1 / NumPy 1.26.4 and CPython 3.11.14 / PyTorch 2.0.1 / NumPy 1.26.4, where test_jvp_torch_compile is skipped by its existing version check.
• Repository-pinned Ruff 0.16.2 format/check across roma, test and examples pass.
• git diff --check passes.

CUDA was not run locally.

[rbregier]

Thanks @poketopa! Rebased onto the current master (resolving a conflict with new tests) and pushed as be34e99, so closing this PR. Your authorship is preserved in the commit.