Best way to pass small inputs as __constant__

[johnpjf]

I have a kernel that does something like multiple an array of 3D vectors by a 4x4 matrix (and then some other stuff).
The kernel signature is like this:

@wp.kernel
def matmul(
    mat: wp.mat44,
    vertices: wp.array2d(dtype=wp.vec4),
    output: wp.array2d(dtype=wp.vec4)):
  i, j = wp.tid()
  output[i, j] = mat @ vertices[i, j]

Note that mat is changing every call.
I don't quite know how to call this signature though, this doesn't work:

    vertices = jnp.ones((100, 100, 4), dtype=jnp.float32)
    mat = jnp.ones((4, 4), dtype=jnp.float32)
    output = jnp.zeros_like(vertices)
    wp_vertices = wp.from_jax(vertices, dtype=wp.vec4)
    wp_mat = wp.from_jax(mat, dtype=wp.mat44)
    wp_output = wp.from_jax(output, dtype=wp.vec4)
    wp.launch(
        kernel=matmul,
        dim=vertices.shape[:2],
        inputs=[
            wp_vertices,
            wp_mat,
            wp_output,
        ],
        device='cuda',
    )
    wp.synchronize()

Can a kernel input be a wp.mat44 like this? If so how do I call it? Also, is there a way to make the wp.mat44 be in __constant__ CUDA RAM?