`dr.diag` no longer accepts `mi.Vector3f`

[dvicini]

The following code used to create a diagonal matrix:

import mitsuba as mi
import drjit as dr 
mi.set_variant('llvm_ad_rgb')
dr.diag(mi.Vector3f(1,2,3))

but now raises Exception: drjit.diag(): unsupported type!.

It works if I use the raw Dr.Jit array type:

import drjit as dr 
dr.diag(dr.llvm.ad.Array3f(1,2,3))

So the issue seems to somehow relate to the way Mitsuba exposes Vector3f, hence I am opening this here and not in the Dr.Jit repo.

[dvicini]

I think the issue is the logic in Dr.Jit's diag function:

    elif is_array_v(arg) and 'Array' in tp.__name__:
        import sys
        mat_tp = getattr(sys.modules[tp.__module__], tp.__name__.replace('Array', 'Matrix'), None)

This does not seem to work for mi.Vector3f, since the type name here does not contain Array

[wjakob]

The more stable way to fix this would be to add a type trait to drjit/src/python/traits.cpp, which produces a matrix type from a vector type of a given size.

[njroussel]

Fixed in https://github.com/mitsuba-renderer/drjit/pull/254