GridTools/gt4py

Maximum does not work for field_operator that are called by another field_operator

huppd opened this issue · 2 comments

huppd commented

This problem only occurs only in the embedded execution with native python.
So the following does work:

@field_operator
def _hflx_limiter_mo_stencil_03_min_max(
    z_tracer_max: Field[[CellDim, KDim], float],
) -> [Field[[CellDim, KDim], float]:
    z_max = maximum(
        max_over(z_tracer_max(C2E2C), axis=C2E2CDim), z_tracer_max
    )
    return z_max

@field_operator
def _hflx_limiter_mo_stencil_03a(
    z_mflx_anti_in: Field[[CellDim, KDim], float],
    z_tracer_new_low: Field[[CellDim, KDim], float],
    z_max: Field[[CellDim, KDim], float],
) -> Field[[CellDim, KDim], float]:
    r_p = (z_max - z_tracer_new_low) / (z_mflx_anti_in)
    return r_p,

@program
def hflx_limiter_mo_stencil_03(
    z_tracer_max: Field[[CellDim, KDim], float],
    z_max: Field[[CellDim, KDim], float],
    z_mflx_anti_in: Field[[CellDim, KDim], float],
    z_tracer_new_low: Field[[CellDim, KDim], float],
    r_p: Field[[CellDim, KDim], float],
):

    _hflx_limiter_mo_stencil_03_min_max(
        z_tracer_max,  out= z_max
    )
    _hflx_limiter_mo_stencil_03a(
        z_mflx_anti_in,
        z_tracer_new_low,
        z_max,
        out=r_p,
    )

the following doesnot

@field_operator
def _hflx_limiter_mo_stencil_03_min_max(
    z_tracer_max: Field[[CellDim, KDim], float],
) -> [Field[[CellDim, KDim], float]:
    z_max = maximum(
        max_over(z_tracer_max(C2E2C), axis=C2E2CDim), z_tracer_max
    )
    return z_max

@field_operator
def _hflx_limiter_mo_stencil_03a(
    z_mflx_anti_in: Field[[CellDim, KDim], float],
    z_tracer_new_low: Field[[CellDim, KDim], float],
    z_max: Field[[CellDim, KDim], float],
) -> Field[[CellDim, KDim], float]:
    r_p = (z_max - z_tracer_new_low) / (z_mflx_anti_in)
    return r_p,

@field_operator
def _hflx_limiter_mo_stencil_03(
    z_tracer_max: Field[[CellDim, KDim], float],
    z_mflx_anti_in: Field[[CellDim, KDim], float],
    z_tracer_new_low: Field[[CellDim, KDim], float],
) -> tuple[Field[[CellDim, KDim], float], Field[[CellDim, KDim], float]]:


    z_max = _hflx_limiter_mo_stencil_03_min_max( z_tracer_max)

    r_p =_hflx_limiter_mo_stencil_03a(
        z_mflx_anti_in,
        z_tracer_new_low,
        z_max,
    )
    return r_p

@program
def hflx_limiter_mo_stencil_03(
    z_tracer_max: Field[[CellDim, KDim], float],
    z_mflx_anti_in: Field[[CellDim, KDim], float],
    z_tracer_new_low: Field[[CellDim, KDim], float],
    r_p: Field[[CellDim, KDim], float],
):

    _hflx_limiter_mo_stencil_03(
        z_tracer_max,
        z_mflx_anti_in,
        z_tracer_new_low,
        out=r_p,
    )

It fails with the error:

FAILED advection/tests/test_hflx_limiter_mo_stencil_03.py::test_hflx_diffusion_mo_stencil_03 - NotImplementedError: Missing type definition for builtin 'maximum'

I cannot reproduce now. Can you confirm that it's fixed with the latest GT4Py?

huppd commented

Now, I can't reproduce it either, so I can confirm that it got fixed. Thanks a lot :)