Problems when indexing inside an UDF

[martaiborra]

Inside an UDF, the elements from the out array cannot be referenced with a constant. The following example fails:

import numpy as np
import iarray as ia
from iarray.udf import jit, Array, float64

@jit()
def udf_max(
    out: Array(float64, 2),
    x: Array(float64, 2)) -> int:

    n_rows = x.window_shape[0]
    n_cols = x.window_shape[1]
    for i in range(n_rows):
        for j in range(n_cols):
            out[i, 0] = n_rows

    return 0

shape = [20, 30]
chunks = [10, 15]
blocks = [5, 5]
dtype = np.float64

iarr = ia.arange(shape=shape, chunks=chunks, blocks=blocks, dtype=dtype)

expr = ia.expr_from_udf(udf_max, [iarr])
out = expr.eval()

ValueError: Operands must be the same type, got (i64, i32)

[martaiborra]

The same happens if its a variable:

import numpy as np
import iarray as ia
from iarray.udf import jit, Array, float64

@jit()
def udf_max(
    out: Array(float64, 2),
    x: Array(float64, 2)) -> int:

    n_rows = x.window_shape[0]
    n_cols = x.window_shape[1]
    cols = 0
    for i in range(n_rows):
        for j in range(n_cols):
                out[i, cols] = x[i, j]

    return 0

shape = [20, 30]
chunks = [10, 15]
blocks = [5, 5]
dtype = np.float64

iarr = ia.arange(shape=shape, chunks=chunks, blocks=blocks, dtype=dtype)

expr = ia.expr_from_udf(udf_max, [iarr])
out = expr.eval()

ValueError: Operands must be the same type, got (i64, i32)

[jdavid]

Fixed in https://github.com/inaos/iron-array-python/pull/143