How to not return a readonly array in numba objmode?

user_19240582 0
Viewed 32

Problem: I have a NumbaPerformanceWarning about the use of non-contiguous arrays when using objmode in nopython mode, as it seems to return a view of the objects.

Bellow is a small code to reproduce my problem:

import numpy as np
from numba import njit, objmode
from scipy.linalg import expm

@njit
def foo(A, b):
    # created contiguous
    G_t = np.zeros(shape=A.shape, dtype=np.float64)

    with objmode(G_t="float64[:, :]"):
        G_t += expm(A)
        # When uncommented -> still a NumbaPerformanceWarning
        # G_t = np.ascontiguousarray(G_t)

    # When uncommented -> no more NumbaPerformanceWarning
    # G_t = np.ascontiguousarray(G_t)
    return G_t @ b

N = 4
A = np.random.random(size=(N, N))
b = np.random.random(size=(N))
foo(A, b)
# NumbaPerformanceWarning: '@' is faster on contiguous arrays, called on (array(float64, 2d, A), readonly array(float64, 1d, C))
# return G_t @ b

Question: Is my use of objmode wrong? Is there a clean way to prevent this performance warning?

Thanks !

numba python-3.x
1 Answers

float64[:, :] means a view that is not mandatory contiguous while float64[:, ::1] is always contiguous. Thus, using the second one make the code faster and removes the warning here. Note that doing the += in the objmode section is not the most efficient solution. You extract the array created in the section and to the += after that.

with objmode(tmp="float64[:, ::1]"):
    tmp = expm(A)
G_t += tmp

This solution is 20% faster than just adding ::1 (which nearly do not improves performance because Numpy do not use Numba typing information).

user_129395570
Related