提交 426e0353 authored 作者: Adrian Seyboldt's avatar Adrian Seyboldt 提交者: Ricardo Vieira

perf(numba): Avoid casting arrays if not necessary

上级 6d3c7568
...@@ -746,7 +746,13 @@ def numba_funcify_SpecifyShape(op, node, **kwargs): ...@@ -746,7 +746,13 @@ def numba_funcify_SpecifyShape(op, node, **kwargs):
def int_to_float_fn(inputs, out_dtype): def int_to_float_fn(inputs, out_dtype):
"""Create a Numba function that converts integer and boolean ``ndarray``s to floats.""" """Create a Numba function that converts integer and boolean ``ndarray``s to floats."""
if any(i.type.numpy_dtype.kind in "ib" for i in inputs): if all(input.type.numpy_dtype == np.dtype(out_dtype) for input in inputs):
@numba_njit
def inputs_cast(x):
return x
elif any(i.type.numpy_dtype.kind in "ib" for i in inputs):
args_dtype = np.dtype(f"f{out_dtype.itemsize}") args_dtype = np.dtype(f"f{out_dtype.itemsize}")
@numba_njit @numba_njit
......
Markdown 格式
0%
您添加了 0 到此讨论。请谨慎行事。
请先完成此评论的编辑!
注册 或者 后发表评论