Problems with the window_shape in UDF

[martaiborra]

In a UDF, the window_shape does not coincide with the block shape nor the chunk shape. This can be seen in the following example:

import numpy as np
import math
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, j] = n_cols

    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()

The chunk shape is [10, 15], but in the columns, the window_shape is 16 for the first one and 14 for the other one and it should be 15. For the rows it is 8 for the first ones and 4 for the others and it should be 10.

[FrancescAlted]

@martaiborra Please attach the output so that it is easier to assess from the "casual" visitor.

[martaiborra]

The output array obtained when out[i, j] = n_cols is:

[[16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 14. 14.
  14. 14. 14. 14. 14. 14. 14. 14. 14. 14. 14. 14.]
 [16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 14. 14.
  14. 14. 14. 14. 14. 14. 14. 14. 14. 14. 14. 14.]
 [16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 14. 14.
  14. 14. 14. 14. 14. 14. 14. 14. 14. 14. 14. 14.]
 [16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 14. 14.
  14. 14. 14. 14. 14. 14. 14. 14. 14. 14. 14. 14.]
 [16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 14. 14.
  14. 14. 14. 14. 14. 14. 14. 14. 14. 14. 14. 14.]
 [16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 14. 14.
  14. 14. 14. 14. 14. 14. 14. 14. 14. 14. 14. 14.]
 [16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 14. 14.
  14. 14. 14. 14. 14. 14. 14. 14. 14. 14. 14. 14.]
 [16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 14. 14.
  14. 14. 14. 14. 14. 14. 14. 14. 14. 14. 14. 14.]
 [16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 14. 14.
  14. 14. 14. 14. 14. 14. 14. 14. 14. 14. 14. 14.]
 [16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 14. 14.
  14. 14. 14. 14. 14. 14. 14. 14. 14. 14. 14. 14.]
 [16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 14. 14.
  14. 14. 14. 14. 14. 14. 14. 14. 14. 14. 14. 14.]
 [16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 14. 14.
  14. 14. 14. 14. 14. 14. 14. 14. 14. 14. 14. 14.]
 [16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 14. 14.
  14. 14. 14. 14. 14. 14. 14. 14. 14. 14. 14. 14.]
 [16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 14. 14.
  14. 14. 14. 14. 14. 14. 14. 14. 14. 14. 14. 14.]
 [16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 14. 14.
  14. 14. 14. 14. 14. 14. 14. 14. 14. 14. 14. 14.]
 [16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 14. 14.
  14. 14. 14. 14. 14. 14. 14. 14. 14. 14. 14. 14.]
 [16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 14. 14.
  14. 14. 14. 14. 14. 14. 14. 14. 14. 14. 14. 14.]
 [16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 14. 14.
  14. 14. 14. 14. 14. 14. 14. 14. 14. 14. 14. 14.]
 [16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 14. 14.
  14. 14. 14. 14. 14. 14. 14. 14. 14. 14. 14. 14.]
 [16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 14. 14.
  14. 14. 14. 14. 14. 14. 14. 14. 14. 14. 14. 14.]]

And the output array obtained when out[i, j] = n_rows is:

[[8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8.
  8. 8. 8. 8. 8. 8.]
 [8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8.
  8. 8. 8. 8. 8. 8.]
 [8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8.
  8. 8. 8. 8. 8. 8.]
 [8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8.
  8. 8. 8. 8. 8. 8.]
 [8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8.
  8. 8. 8. 8. 8. 8.]
 [8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8.
  8. 8. 8. 8. 8. 8.]
 [8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8.
  8. 8. 8. 8. 8. 8.]
 [8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8.
  8. 8. 8. 8. 8. 8.]
 [8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8.
  8. 8. 8. 8. 8. 8.]
 [8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8.
  8. 8. 8. 8. 8. 8.]
 [8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8.
  8. 8. 8. 8. 8. 8.]
 [8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8.
  8. 8. 8. 8. 8. 8.]
 [8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8.
  8. 8. 8. 8. 8. 8.]
 [8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8.
  8. 8. 8. 8. 8. 8.]
 [8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8.
  8. 8. 8. 8. 8. 8.]
 [8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8. 8.
  8. 8. 8. 8. 8. 8.]
 [4. 4. 4. 4. 4. 4. 4. 4. 4. 4. 4. 4. 4. 4. 4. 4. 4. 4. 4. 4. 4. 4. 4. 4.
  4. 4. 4. 4. 4. 4.]
 [4. 4. 4. 4. 4. 4. 4. 4. 4. 4. 4. 4. 4. 4. 4. 4. 4. 4. 4. 4. 4. 4. 4. 4.
  4. 4. 4. 4. 4. 4.]
 [4. 4. 4. 4. 4. 4. 4. 4. 4. 4. 4. 4. 4. 4. 4. 4. 4. 4. 4. 4. 4. 4. 4. 4.
  4. 4. 4. 4. 4. 4.]
 [4. 4. 4. 4. 4. 4. 4. 4. 4. 4. 4. 4. 4. 4. 4. 4. 4. 4. 4. 4. 4. 4. 4. 4.
  4. 4. 4. 4. 4. 4.]]

[jdavid]

I think window_shape should be the block shape and not the chunk shape, so in the example above it should be 5x5.

If the chunk/block shape is defined explicitly it works:

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

iarr = ia.arange(shape=shape, chunks=chunks, blocks=blocks, dtype=dtype)
cfg = ia.Config(chunks=chunks, blocks=blocks)
expr = ia.expr_from_udf(udf_max, [iarr], shape=shape, cfg=cfg)
out = expr.eval()

Then the output is correct:

[[5. 5. 5. 5. 5. 5. 5. 5. 5. 5. 5. 5. 5. 5. 5. 5. 5. 5. 5. 5. 5. 5. 5. 5.
  5. 5. 5. 5. 5. 5.]
 [5. 5. 5. 5. 5. 5. 5. 5. 5. 5. 5. 5. 5. 5. 5. 5. 5. 5. 5. 5. 5. 5. 5. 5.
  5. 5. 5. 5. 5. 5.]
...

[FrancescAlted]

That's a good point. Apparently, the chunks/blocks configuration does not propagate correctly to the output for the UDF. That's strange because the expression machinery should just copy the chunks/blocks for the output from the input, so something is not working as it should. @martaiborra can you have a look into this?

[martaiborra]

Fixed in https://github.com/inaos/iron-array-python/commit/4f4028606ffdae5c922dd3d1a9c71db979adcea9