[BUG] `wp.svd2` returns NaN in `U` and `V`

[akrivx]

Bug Description

There seems to be an issue with the wp.svd2 function, which consistently returns NaN values in the U and V matrices. Even for a simple diagonal input, U and V are filled with NaNs, although singular values appear correctly computed.

Minimal Reproducible Example

import warp as wp
wp.init()

@wp.kernel
def test_svd2(U: wp.array(dtype=wp.mat22), s: wp.array(dtype=wp.vec2), V: wp.array(dtype=wp.mat22)):
    U2 = wp.mat22()
    V2 = wp.mat22()
    s2 = wp.vec2()
    wp.svd2(wp.diag(wp.vec2(2.0)), U2, s2, V2)
    U[0] = U2
    s[0] = s2
    V[0] = V2

U = wp.zeros(1, dtype=wp.mat22)
V = wp.zeros(1, dtype=wp.mat22)
s = wp.zeros(1, dtype=wp.vec2)

wp.launch(kernel=test_svd2, dim=1, inputs=[U, s, V])

print("U:", U)
print("V:", V)
print("s:", s)

Observed Output

U: [[[nan nan]
     [nan nan]]]
V: [[[nan nan]
     [nan nan]]]
s: [[2. 2.]]

Expected Behaviour

U and V should contain valid orthogonal matrices (no NaNs) resulting from the decomposition of a simple diagonal matrix.

System Information

• NVIDIA Warp Version: 1.6.2
• CUDA Version: CUDA Toolkit 12.8, Driver 12.9
• GPU Model: NVIDIA GeForce RTX 5080
• Python Version: 3.12.

[shi-eric]

@gdaviet added improvements to wp.svd2 which should address the issue reported here. It's available in the main branch already and will be available in future nightly builds. See 6f0b1c25b28609ee851592bc0ecc8d8b34cb33e5