[Bug Report] Unsupported `transform_decompose` and `transform_compose` in CUDA and LLVM backend

[Chaphlagical]

Hi, I've come across a issue when performing transform_decompose and transform_compose with mi.Matrix4f in CUDA and LLVM backend:

import mitsuba as mi
import drjit as dr

mi.set_variant("cuda_ad_rgb")
# mi.set_variant("llvm_ad_rgb")

m = dr.identity(mi.Matrix4f)
s, q, t = dr.transform_decompose(m)
dr.transform_compose(s, q, t)

I get the following error:

Float constructor expects: arbitrarily many values of type "float", a matching list/tuple, or a NumPy/PyTorch/TF/Jax array.
  File "D:\Workspace\mitsuba\build\Release\python\drjit\detail.py", line 208, in array_init
    self.set_entry_(i, value_type(o[i]))

I found the issue is caused by _dr.detail.array_name in matrix.py where the functions are defined.

def transform_decompose(a, it=10):
    '''
    transform_decompose(arg, it=10)
    Performs a polar decomposition of a non-perspective 4x4 homogeneous
    coordinate matrix and returns a tuple of

    1. A positive definite 3x3 matrix containing an inhomogeneous scaling operation
    2. A rotation quaternion
    3. A 3D translation vector

    This representation is helpful when animating keyframe animations.

    Args:
        arg (drjit.ArrayBase): A Dr.Jit matrix type

        it (int): Number of iterations to be taken by the polar decomposition algorithm.

    Returns:
        tuple: The tuple containing the scaling matrix, rotation quaternion and 3D translation vector.
    '''
    if not _dr.is_matrix_v(a):
        raise Exception('Unsupported type!')

    name = _dr.detail.array_name('Array', a.Type, [3], a.IsScalar)
    module = _modules.get(a.__module__)
    Array3f = getattr(module, name)

    name = _dr.detail.array_name('Matrix', a.Type, (3, 3), a.IsScalar)
    Matrix3f = getattr(module, name)

    Q, P = polar_decomp(Matrix3f(a), it)

    sign_q = det(Q)
    Q = _dr.mulsign(Q, sign_q)
    P = _dr.mulsign(P, sign_q)

    return P, matrix_to_quat(Q), Array3f(a[3][0], a[3][1], a[3][2])

def transform_compose(s, q, t, /):
    '''
    transform_compose(S, Q, T, /)
    This function composes a 4x4 homogeneous coordinate transformation from the
    given scale, rotation, and translation. It performs the reverse of
    :py:func:`transform_decompose`.

    Args:
        S (drjit.ArrayBase): A Dr.Jit matrix type representing the scaling part
        Q (drjit.ArrayBase): A Dr.Jit quaternion type representing the rotation part
        T (drjit.ArrayBase): A 3D Dr.Jit array type representing the translation part

    Returns:
        drjit.ArrayBase: The Dr.Jit matrix resulting from the composition described above.
    '''
    if not _dr.is_matrix_v(s) or not _dr.is_quaternion_v(q):
        raise Exception('Unsupported type!')

    name = _dr.detail.array_name('Matrix', q.Type, (4, 4), q.IsScalar)
    module = _modules.get(q.__module__)
    Matrix4f = getattr(module, name)

    result = Matrix4f(quat_to_matrix(q, 3) @ s)
    result[3] = [t[0], t[1], t[2], 1.0]

    return result

When input type is drjit.cuda(/llvm).ad.Float, _dr.detail.array_name will always return drjit.cuda(/llvm).ad.Float and leading failure in the following computation. Currently I just add a new dimension so that _dr.detail.array_name can capture the correct type:

#...
# name = _dr.detail.array_name('Array', a.Type, [3], a.IsScalar)
name = _dr.detail.array_name('Array', a.Type, [3, 3], a.IsScalar)
# ...
# name = _dr.detail.array_name('Matrix', a.Type, (3, 3), a.IsScalar)
name = _dr.detail.array_name('Matrix', a.Type, (3, 3, 3), a.IsScalar)
# ...

I wonder if there is a better solution.

Best, Chen

[Chaphlagical]

Thank you very much!