[cartesian] Numpy backend lacks non-literal accessor for data dimensions

[FlorianDeconinck]

The numpy backend has a limitation on literal accessor not shared by other backends.
Consider the following:

import gt4py.cartesian.gtscript as gtscript
from gt4py.cartesian.gtscript import (
    computation,
    interval,
    PARALLEL,
)
import numpy as np

THE_NUMBER_5 = 5
FloatField_NModes = gtscript.Field[gtscript.IJK, (np.float32, (THE_NUMBER_5))]
FloatField = gtscript.Field[gtscript.IJK, np.float32]

@gtscript.stencil(backend="gt:cpu_kfirst")
def ddim_read_in_loop(
        an_array_of_ones: FloatField_NModes,
        out_field: FloatField):
    with computation(PARALLEL), interval(...):
        n_iter = 0
        agg = 0
        while n_iter < THE_NUMBER_5:
            agg = agg + an_array_of_ones[0, 0, 0][n_iter]
            n_iter = n_iter + 1
        out_field = agg

the_ones_field = np.ones((2, 2, 2, THE_NUMBER_5), dtype=np.float32)
out_field = np.zeros((2, 2, 2), dtype=np.float32)

ddim_read_in_loop(the_ones_field, out_field)
print(out_field)

We expect n_iter to be some kind of local scalar variable. But numpy is trying to use as much inner parallelism and masking as possible leading to this code giving a

  File "/home/mad/work/ndsl/external/gt4py/src/gt4py/cartesian/gtc/numpy/npir.py", line 141, in data_indices_are_scalar
    raise ValueError("Data indices must be scalars")
ValueError: Data indices must be scalars

Meanwhile, running this code under a gt:Xor dace:X backend gives the expected results (with n_iter a local scalar`).

PS: this was unearth looking at the Aer Activation physics code for NASA's GEOS model

[egparedes]

Thanks for reporting @FlorianDeconinck . Are you planning to also work in a fix or is it not high priority for you at the moment?

[FlorianDeconinck]

Unclear. We need a fix for sure but we have a mounting amount of physics pattern difficult to write so other feature requests might take precedence. I'd say we are a few weeks away from deciding where we put our efforts

[FlorianDeconinck]

So it seems this is going to be a slightly bigger work than a simple bug fix.

The numpy backends tries to use broadcasting as much as possible, which means for example that local scalar access will be considered as vectors, e.g.:

class LocalScalarAccess(VectorLValue):
    name: eve.Coerced[eve.SymbolRef]
    kind: common.ExprKind = common.ExprKind.FIELD

We will have to refactor that first.

[FlorianDeconinck]

Probable solution used for a POC of another feature:

            # The code as been modified to generate B[i:I,j:J,index]
            # which is correct but will break because you can't brodcast 2D in 3D
            # e.g. you can't write (for I, J, K the dimensions)
            #   > (I,J,K) = (I,J)
            # we need to use push the 2D into 3D realm by re-establishing a 3rd axis, e.g:
            #   > A[i:I, j:J, k:K] = B[i:I, j:J, index][..., np.newaxis]
            # This isn't true for FieldSlice which would lead to
            # a (I, J, K) = (I, J, 1), which is structurally equal to the above extra axis add.