Início / coding / Perguntas / 79575363
Accepted
Raven
Asked: 2025-04-15 22:09:27 +0800 CST

como forçar numba a retornar um tipo numpy?

  • 772

Acho esse comportamento bastante contraintuitivo, embora eu suponha que haja uma razão para isso - o numba converte automaticamente meus tipos inteiros numpy diretamente em um int python:

import numba as nb
import numpy as np 

print(f"Numba version: {nb.__version__}")  # 0.59.0
print(f"NumPy version: {np.__version__}")  # 1.23.5

# Explicitly define the signature
sig = nb.uint32(nb.uint32, nb.uint32)

@nb.njit(sig, cache=False)
def test_fn(a, b):
    return a * b

res = test_fn(2, 10)
print(f"Result value: {res}")  # returns 20
print(f"Result type: {type(res)}")  # returns <class 'int'>

Este é um problema porque estou usando o retorno como uma entrada em outra função njit, então recebo um aviso de conversão (e também faço conversões desnecessárias entre as funções njit)

Existe alguma maneira de forçar o numba a me dar np.uint32um resultado?

--- EDITAR ---

Este é o melhor que consegui fazer sozinho, mas me recuso a acreditar que esta seja a melhor implementação que existe:

# we manually define a return record and pass it as a parameter
res_type = np.dtype([('res', np.uint32)])
sig = nb.void(nb.uint32, nb.uint32, nb.from_dtype(res_type))

@nb.njit(sig, cache=False)
def test_fn(a:np.uint32, b:np.uint32, res: res_type):
    res['res'] = a * b

# Call with Python ints (Numba should coerce based on signature)
res = np.recarray(1, dtype=res_type)[0]
res_py_in = test_fn(2, 10, res)
print(f"\nCalled with Python ints:")
print(f"Result value: {res['res']}")  # 20
print(f"Result type: {type(res['res'])}")  # <class 'numpy.uint32'>
python

2 respostas

  1. Best Answer

    Você pode gostar desta, pois é limpa, menos prolixa, interoperável com outras funções Numba e o valor de retorno se comporta como um escalar, mas tem o dtype correto. Ideal para uso geral.

    import numba as nb
import numpy as np

@nb.njit
def test_fn(a, b):
    res = np.empty((), dtype=np.uint32)
    res[...] = a * b
    return res

# Test
a = np.uint32(2)
b = np.uint32(10)
result = test_fn(a, b)
print("Result value:", result)
print("Type of result:", type(result))
print("Dtype of result:", result.dtype)
print("As Python int (optional):", result.item())

    Outro (evita alocação de memória, é mais eficiente em loops apertados e melhor para cenários de desempenho crítico):

    res_type = np.dtype([('res', np.uint32)])
sig = nb.void(nb.uint32, nb.uint32, nb.from_dtype(res_type))

@nb.njit(sig)
def test_fn(a, b, out):
    out['res'] = a * b

res = np.recarray(1, dtype=res_type)[0]
test_fn(np.uint32(2), np.uint32(10), res)
print("Result value:", res['res'])
print("Type of result:", type(res['res']))
print("As Python int (optional):", int(res['res']))

    Saída:

    D:\python>python test.py
Result value: 20
Type of result: <class 'numpy.uint32'>
As Python int (optional): 20
    • 2

  2. Por que você simplesmente não retorna np.uint32(a*b):

    @nb.njit(nb.uint32(nb.uint32, nb.uint32))
def func(a, b):
    return np.uint32(a * b)

    É mais rápido e legível que as outras soluções:

    import numba as nb
import numpy as np

@nb.njit(nb.types.Array(nb.uint32, 0, "C")(nb.uint32, nb.uint32))
def test_fn(a, b):
    res = np.empty((), dtype=np.uint32)
    res[...] = a * b
    return res

res_type = np.dtype([('res', np.uint32)])
sig = nb.void(nb.uint32, nb.uint32, nb.from_dtype(res_type))

@nb.njit(sig)
def test_fn2(a, b, out):
    out['res'] = a * b

res = np.recarray(1, dtype=res_type)[0]
test_fn2(np.uint32(2), np.uint32(10), res)

a = np.uint32(2)
b = np.uint32(10)


%timeit test_fn(a, b)
%timeit test_fn2(a, b, res)
%timeit func(a, b)

    Saída:

    339 ns ± 4.67 ns per loop (mean ± std. dev. of 7 runs, 1,000,000 loops each)
426 ns ± 1.01 ns per loop (mean ± std. dev. of 7 runs, 1,000,000 loops each)
126 ns ± 0.111 ns per loop (mean ± std. dev. of 7 runs, 10,000,000 loops each)

    N = int(1e7)
@nb.njit
def _test_fn(a, b):
    out = np.empty((N,), dtype=np.uint32)
    for i in range(N):
        out[i] = test_fn(a, b).item()
    return out

@nb.njit
def _test_fn2(a, b, res):
    out = np.empty((N,), dtype=np.uint32)
    for i in range(N):
        test_fn2(a, b, res)
        out[i] = res['res']
    return out

@nb.njit
def _func(a, b):
    out = np.empty((N,), dtype=np.uint32)
    for i in range(N):
        out[i] = func(a, b)
    return out

_test_fn(a, b)
_test_fn2(a, b, res)
_func(a, b)

%timeit _test_fn(a, b)
%timeit _test_fn2(a, b, res)
%timeit _func(a, b)

    Saída:

    254 ms ± 508 μs per loop (mean ± std. dev. of 7 runs, 1 loop each)
3.44 ms ± 40.4 μs per loop (mean ± std. dev. of 7 runs, 100 loops each)
3.37 ms ± 19.9 μs per loop (mean ± std. dev. of 7 runs, 100 loops each)
    • 1

relate perguntas