4

I am looking for any alternate way to write

arr_1d = np.zeros(d ** n)
arr_nd = arr_1d.reshape((d,) * n)

with a @numba.jit(nopython=True) decorator. (The rank n is not fixed so writing (d,...,d) is not a solution)

As an example:

from numba import njit, objmode

def tensor(series: np.ndarray, dim: int):
    n = int(np.log(len(series)) / np.log(dim))
    new_shape = (dim,) * n
    # new_shape = [dim] * n
    # new_shape = tuple([dim] * n)
    # with objmode(new_shape="UniTuple(dtype=np.int8, count="+str(n)+")"):
    #     new_shape = (dim,) * n
    return series.reshape(new_shape)

d = 2
n = 3
series = np.zeros(shape=d**n)

print(tensor(series=series, dim=d))
# [[[0. 0.]
#   [0. 0.]]
# 
#  [[0. 0.]
#   [0. 0.]]]

print(njit(tensor)(series=series, dim=d))
# Error (see below)

I tried different ways of writing the new_shape tuple:

The multiplication of tuples seems not to be implemented:

new_shape = (dim,) * n
# TypingError: No implementation of function Function(<built-in function mul>) found for signature

The multiplication of lists is implemented, but a list does not work in the reshape method:

new_shape = [dim] * n
# TypingError: Invalid use of BoundFunction(array.reshape for array(float64, 1d, C)) with parameters (list(int64)<iv=None>)

As stated in the documentation, but worth a shot, tuple() is not implemented in numba:

new_shape = tuple([dim] * n)
# TypingError: No implementation of function Function(<class 'tuple'>) found for signature:

So I tried objmode but I don't understand it enough, so I can not make it work...

with objmode(new_shape="UniTuple(np.int8,"+str(n)+")"):
    new_shape = (dim,) * n
# CompilerError: Error handling objmode argument 'new_shape'. The value must be a compile-time constant either as a non-local variable or a getattr expression that refers to a Numba type.

There might not be a solution yet, but if I missed something, I'd be delighted to receive any advice !

Thanks

Louis-Amand
  • 101
  • 7
  • 3
    I think it's not even possible to have varying array shapes as a return value. [This answer](https://stackoverflow.com/a/70598981/14277722) might be helpful, but it's about a different topic regarding `tuples`. – Michael Szczesny May 31 '22 at 14:42
  • 1
    You can also make `new_shape` by `np.ones(n, dtype = np.int64) * dim` but that doesn't work either, but it's a different error at least (`Invalid use of BoundFunction(array.reshape for array(float64, 1d, C)) with parameters (array(int64, 1d, C))`) – Daniel F May 31 '22 at 14:43
  • 1
    Generally you can't do this within the Numba function. The array dimensionality has to be known at compile time and can't be changed at runtime. But there is a possible workaround. You can use a generator function to get functions for all needed array dims. In this answer https://stackoverflow.com/a/59356461/4045774 I used a generator function to get optimized matrix-matrix multiplication, depended on array shapes. You probably can do quite the same for different array-dims. Of course another possible solution would be to do the reshape outside of Numba. – max9111 May 31 '22 at 15:03

0 Answers0