dynamic input shapes (grid) supported?

[nlgranger]

I'm getting custom_einsum_triton() got multiple values for argument 'grid' as soon as I submit an input with a different shape and therefore grid size. Are variable grid sizes supported?

@triton.jit
def custom_einsum_triton(
    xyz_ptr,
    xyz_stride_i,
    xyz_stride_j,
    sigma_inv_ptr,
    sigma_inv_stride_i,
    sigma_inv_stride_j,
    sigma_inv_stride_k,
    out_ptr,
    out_stride_i,
    N,
    DIM: tl.constexpr,
    BLOCK_SIZE: tl.constexpr,
):
    ...

BLOCK_SIZE = 1024
grid = lambda meta: (triton.cdiv(out.shape[0], meta["BLOCK_SIZE"]),)

custom_einsum_triton[grid](
    xyz, *xyz.stride(),
    sigma_inv, *sigma_inv.stride(),
    out, *out.stride(),
    xyz.shape[0], xyz.shape[1], BLOCK_SIZE)

[nlgranger]

Scratch that: it is.
My error stemmed from feeding tensors with extra dims, which I did not implement. I guess the error message could be improved, but the mistake was on my side.