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.
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