matmul-ia-excerpt3.py

from iarray.udf import jit, Array, float64, int64

# Create second array via an UDF
@jit()
def tri(out: Array(float64, 2), x: Array(float64, 2), k: int64) -> int64:
    n = out.window_shape[0]
    m = out.window_shape[1]
    row_start = out.window_start[0]
    col_start = out.window_start[1]
    for i in range(n):
        for j in range(m):
            if (row_start + i) >= (col_start + j - k):
                out[i, j] = 1
            else:
                out[i, j] = 0
    return 0

# For the UDF input we don't need actual values, so an empty array is enough
ia_in = ia.empty(bmshape, chunks=bmchunks, blocks=bmblocks)
expr = ia.expr_from_udf(tri, [ia_in], [1])
bm = expr.eval()