Numba.vectorize with target='cuda' return type is not respected

import numpy as np
from numba import vectorize, float32

@vectorize([float32(float32)], target='cuda')
def vec_fn(x):
    return max(x, 0.0)

vec_fn(np.array([1.0], dtype="float32")).dtype

output

dtype('float64')

expected

dtype('float32')

Am I missing something?

Without target='cuda' it does the right thing.

Thanks!

Thanks for bringing this up - I can reproduce the behaviour:

import numpy as np
from numba import vectorize, float32


def vec_fn(x):
    return max(x, 0.0)

cpu_vec_fn = vectorize([float32(float32)])(vec_fn)
cuda_vec_fn = vectorize([float32(float32)], target='cuda')(vec_fn)

print(cpu_vec_fn(np.array([1.0], dtype="float32")).dtype)
print(cuda_vec_fn(np.array([1.0], dtype="float32")).dtype)

prints

float32
float64

I’m having a quick look into this now.

Turns out this is Issue #8400 and I had it in the milestone to fix for the 0.57 release. I’ve not yet actually started working on the fix, so I’m hoping it’s just some missing setup for wiring up the return type correctly.

Awesome thanks Graham.

I think I worked out a start on the fix: On target cuda, return type of function may not match type given in signature string · Issue #8400 · numba/numba · GitHub

It could do with a bit more testing than I’ve given it to make it into a PR, but I think the fundamental change is there.

1 Like