Repository: GridTools/gt4py
Pull request #1443 (2024-02-01 23:19:01)

fix[next][dace]: Fix for neighbor reduction with skip values
This PR provides a bugfix for the case of neighbor reductions with lambda function as reduction operation and connectivity table containing skip values. This specific case is verified by `test_tuple_with_local_field_in_reduction_shifted` and the corresponding field operator ITIR looks like:
```
  out ← (
      λ(e, v) → 
        reduce(λ(_fuse_maps_1, _fuse_maps_3, _fuse_maps_4) → _fuse_maps_1 + (_fuse_maps_3 + _fuse_maps_4), 0)(
          neighbors(V2Eₒ, ⟪E2Vₒ, 0ₒ⟫(e)), make_const_list(·⟪E2Vₒ, 0ₒ⟫(v))
        )
    )(__sym_1, __sym_2) @ u⟨ Edge: [0, __out_size_0) ⟩;
```
The lambda function should only accumulate the results for the valid  neighbors. On the contrary, the baseline implementation was using the reduction identity value for the missing neighbors, resulting in invalid result.

The fix consists of producing an array of boolean flags to determine if the neighbor value is valid or not. If not valid, the call to the lambda function is by-passed.

![image](https://github.com/GridTools/gt4py/assets/25990065/4e2bb158-b5bb-4300-a67c-c49b0321ef30)

Code documentation is missing. The intention is to add it after rebase of #1431 (support lowering of lift expressions in neighbor reductions) which will conflict with these changes.


Patch:
diff --git a/src/gt4py/next/iterator/embedded.py b/src/gt4py/next/iterator/embedded.py
index 7e0e0608..011ca4d9 100644
--- a/src/gt4py/next/iterator/embedded.py
+++ b/src/gt4py/next/iterator/embedded.py
@@ -168,7 +168,7 @@ class LocatedField(Protocol):
 
     @property
     @abc.abstractmethod
-    def __gt_domain__(self) -> common.Domain: ...
+    def dims(self) -> tuple[common.Dimension, ...]: ...
 
     # TODO(havogt): define generic Protocol to provide a concrete return type
     @abc.abstractmethod
@@ -176,7 +176,7 @@ class LocatedField(Protocol):
 
     @property
     def __gt_origin__(self) -> tuple[int, ...]:
-        return tuple([0] * len(self.__gt_domain__.dims))
+        return tuple([0] * len(self.dims))
 
 
 @runtime_checkable
@@ -678,18 +678,12 @@ def _is_concrete_position(pos: Position) -> TypeGuard[ConcretePosition]:
 def _get_axes(
     field_or_tuple: LocatedField | tuple,
 ) -> Sequence[common.Dimension]:  # arbitrary nesting of tuples of LocatedField
-    return _get_domain(field_or_tuple).dims
-
-
-def _get_domain(
-    field_or_tuple: LocatedField | tuple,
-) -> common.Domain:  # arbitrary nesting of tuples of LocatedField
     if isinstance(field_or_tuple, tuple):
-        first = _get_domain(field_or_tuple[0])
-        assert all(first == _get_domain(f) for f in field_or_tuple)
+        first = _get_axes(field_or_tuple[0])
+        assert all(first == _get_axes(f) for f in field_or_tuple)
         return first
     else:
-        return field_or_tuple.__gt_domain__
+        return field_or_tuple.dims
 
 
 def _single_vertical_idx(
@@ -900,8 +894,8 @@ class NDArrayLocatedFieldWrapper(MutableLocatedField):
     _ndarrayfield: common.Field
 
     @property
-    def __gt_domain__(self) -> common.Domain:
-        return self._ndarrayfield.__gt_domain__
+    def dims(self) -> tuple[common.Dimension, ...]:
+        return self._ndarrayfield.__gt_domain__.dims
 
     def _translate_named_indices(
         self, _named_indices: NamedFieldIndices
@@ -1452,7 +1446,7 @@ def _tuple_assign(field: tuple | MutableLocatedField, value: Any, named_indices:
 class TupleOfFields(TupleField):
     def __init__(self, data):
         self.data = data
-        self.__gt_domain__ = _get_domain(data)
+        self.dims = _get_axes(data)
 
     def field_getitem(self, named_indices: NamedFieldIndices) -> Any:
         return _build_tuple_result(self.data, named_indices)
diff --git a/src/gt4py/next/program_processors/runners/dace_iterator/__init__.py b/src/gt4py/next/program_processors/runners/dace_iterator/__init__.py
index ed68e66b..2e9a66c4 100644
--- a/src/gt4py/next/program_processors/runners/dace_iterator/__init__.py
+++ b/src/gt4py/next/program_processors/runners/dace_iterator/__init__.py
@@ -299,7 +299,7 @@ def build_sdfg_from_itir(
     for nested_sdfg in sdfg.all_sdfgs_recursive():
         if not nested_sdfg.debuginfo:
             _, frameinfo = warnings.warn(
-                f"{nested_sdfg} does not have debuginfo. Consider adding them in the corresponding nested sdfg."
+                f"{nested_sdfg.label} does not have debuginfo. Consider adding them in the corresponding nested sdfg."
             ), getframeinfo(
                 currentframe()  # type: ignore
             )
diff --git a/src/gt4py/next/program_processors/runners/dace_iterator/itir_to_tasklet.py b/src/gt4py/next/program_processors/runners/dace_iterator/itir_to_tasklet.py
index ba969608..56ffe7e1 100644
--- a/src/gt4py/next/program_processors/runners/dace_iterator/itir_to_tasklet.py
+++ b/src/gt4py/next/program_processors/runners/dace_iterator/itir_to_tasklet.py
@@ -18,11 +18,11 @@ from typing import Any, Callable, Optional, TypeAlias, cast
 
 import dace
 import numpy as np
-from dace.transformation.dataflow import MapFusion
 
 import gt4py.eve.codegen
 from gt4py import eve
 from gt4py.next import Dimension, StridedNeighborOffsetProvider, type_inference as next_typing
+from gt4py.next.common import SKIP_VALUE as neighbor_skip_value
 from gt4py.next.iterator import ir as itir, type_inference as itir_typing
 from gt4py.next.iterator.embedded import NeighborTableOffsetProvider
 from gt4py.next.iterator.ir import FunCall, Lambda
@@ -184,86 +184,93 @@ class Context:
 def builtin_neighbors(
     transformer: "PythonTaskletCodegen", node: itir.Expr, node_args: list[itir.Expr]
 ) -> list[ValueExpr]:
-    di = dace_debuginfo(node, transformer.context.body.debuginfo)
+    sdfg: dace.SDFG = transformer.context.body
+    state: dace.SDFGState = transformer.context.state
+
+    di = dace_debuginfo(node, sdfg.debuginfo)
     offset_literal, data = node_args
     assert isinstance(offset_literal, itir.OffsetLiteral)
     offset_dim = offset_literal.value
     assert isinstance(offset_dim, str)
-    iterator = transformer.visit(data)
-    assert isinstance(iterator, IteratorExpr)
-    field_desc = iterator.field.desc(transformer.context.body)
-
-    field_index = "__field_idx"
     offset_provider = transformer.offset_provider[offset_dim]
-    if isinstance(offset_provider, NeighborTableOffsetProvider):
-        neighbor_check = f"{field_index} >= 0"
-    elif isinstance(offset_provider, StridedNeighborOffsetProvider):
-        neighbor_check = f"{field_index} < {field_desc.shape[offset_provider.neighbor_axis.value]}"
-    else:
-        assert isinstance(offset_provider, Dimension)
+    if not isinstance(offset_provider, NeighborTableOffsetProvider):
         raise NotImplementedError(
-            "Neighbor reductions for cartesian grids not implemented in DaCe backend."
+            "Neighbor reduction only implemented for connectivity based on neighbor tables."
         )
 
-    assert transformer.context.reduce_identity is not None
-
-    sdfg: dace.SDFG = transformer.context.body
-    state: dace.SDFGState = transformer.context.state
+    iterator = transformer.visit(data)
+    assert isinstance(iterator, IteratorExpr)
+    field_desc = iterator.field.desc(transformer.context.body)
+    origin_index_node = iterator.indices[offset_provider.origin_axis.value]
 
-    shifted_dim = offset_provider.origin_axis.value
+    assert transformer.context.reduce_identity is not None
+    assert transformer.context.reduce_identity.dtype == iterator.dtype
 
-    result_name = unique_var_name()
+    # gather the neighbors in a result array dimensioned for `max_neighbors`
+    neighbor_value_var = unique_var_name()
     sdfg.add_array(
-        result_name, dtype=iterator.dtype, shape=(offset_provider.max_neighbors,), transient=True
+        neighbor_value_var,
+        dtype=iterator.dtype,
+        shape=(offset_provider.max_neighbors,),
+        transient=True,
     )
-    result_access = state.add_access(result_name, debuginfo=di)
+    neighbor_value_node = state.add_access(neighbor_value_var, debuginfo=di)
+
+    # allocate scalar to store index for direct addressing of neighbor field
+    neighbor_index_var = unique_var_name()
+    sdfg.add_scalar(neighbor_index_var, _INDEX_DTYPE, transient=True)
+    neighbor_index_node = state.add_access(neighbor_index_var, debuginfo=di)
 
     # generate unique map index name to avoid conflict with other maps inside same state
-    neighbor_index = unique_name("neighbor_idx")
+    neighbor_map_index = unique_name(f"{offset_dim}_neighbor_map_idx")
     me, mx = state.add_map(
-        f"{offset_dim}_neighbors_map",
-        ndrange={neighbor_index: f"0:{offset_provider.max_neighbors}"},
+        f"{offset_dim}_neighbor_map",
+        ndrange={neighbor_map_index: f"0:{offset_provider.max_neighbors}"},
         debuginfo=di,
     )
-    table_name = connectivity_identifier(offset_dim)
-    table_subset = (f"0:{sdfg.arrays[table_name].shape[0]}", neighbor_index)
 
+    table_name = connectivity_identifier(offset_dim)
     shift_tasklet = state.add_tasklet(
         "shift",
-        code="__result = __table[__idx]",
+        code=f"__result = __table[__idx, {neighbor_map_index}]",
         inputs={"__table", "__idx"},
         outputs={"__result"},
         debuginfo=di,
     )
-    data_access_tasklet = state.add_tasklet(
-        "data_access",
-        code=f"__result = __field[{field_index}]"
-        + (
-            f" if {neighbor_check} else {transformer.context.reduce_identity.value}"
-            if offset_provider.has_skip_values
-            else ""
-        ),
-        inputs={"__field", field_index},
-        outputs={"__result"},
-        debuginfo=di,
-    )
-    idx_name = unique_var_name()
-    sdfg.add_scalar(idx_name, _INDEX_DTYPE, transient=True)
     state.add_memlet_path(
         state.add_access(table_name, debuginfo=di),
         me,
         shift_tasklet,
-        memlet=create_memlet_at(table_name, table_subset),
+        memlet=create_memlet_full(table_name, sdfg.arrays[table_name]),
         dst_conn="__table",
     )
     state.add_memlet_path(
-        iterator.indices[shifted_dim],
+        origin_index_node,
         me,
         shift_tasklet,
-        memlet=dace.Memlet(data=iterator.indices[shifted_dim].data, subset="0", debuginfo=di),
+        memlet=dace.Memlet(data=origin_index_node.data, subset="0", debuginfo=di),
         dst_conn="__idx",
     )
-    state.add_edge(shift_tasklet, "__result", data_access_tasklet, field_index, dace.Memlet())
+    state.add_edge(
+        shift_tasklet,
+        "__result",
+        neighbor_index_node,
+        None,
+        dace.Memlet(data=neighbor_index_var, subset="0"),
+    )
+
+    data_access_tasklet = state.add_tasklet(
+        "data_access",
+        code="__data = __field[__idx]"
+        + (
+            f" if __idx != {neighbor_skip_value} else {transformer.context.reduce_identity.value}"
+            if offset_provider.has_skip_values
+            else ""
+        ),
+        inputs={"__field", "__idx"},
+        outputs={"__data"},
+        debuginfo=di,
+    )
     # select full shape only in the neighbor-axis dimension
     field_subset = tuple(
         f"0:{shape}" if dim == offset_provider.neighbor_axis.value else f"i_{dim}"
@@ -276,15 +283,63 @@ def builtin_neighbors(
         memlet=create_memlet_at(iterator.field.data, field_subset),
         dst_conn="__field",
     )
+    state.add_edge(
+        neighbor_index_node,
+        None,
+        data_access_tasklet,
+        "__idx",
+        dace.Memlet(data=neighbor_index_var, subset="0"),
+    )
     state.add_memlet_path(
         data_access_tasklet,
         mx,
-        result_access,
-        memlet=dace.Memlet(data=result_name, subset=neighbor_index, debuginfo=di),
-        src_conn="__result",
+        neighbor_value_node,
+        memlet=dace.Memlet(data=neighbor_value_var, subset=neighbor_map_index, debuginfo=di),
+        src_conn="__data",
     )
 
-    return [ValueExpr(result_access, iterator.dtype)]
+    if not offset_provider.has_skip_values:
+        return [ValueExpr(neighbor_value_node, iterator.dtype)]
+    else:
+        """
+        In case of neighbor tables with skip values, in addition to the array of neighbor values this function also
+        returns an array of booleans to indicate if the neighbor value is present or not. This node is only used
+        for neighbor reductions with lambda functions, a very specific case. For single input neighbor reductions,
+        the regular case, this node will be removed by the simplify pass.
+        """
+        neighbor_valid_var = unique_var_name()
+        sdfg.add_array(
+            neighbor_valid_var,
+            dtype=dace.dtypes.bool,
+            shape=(offset_provider.max_neighbors,),
+            transient=True,
+        )
+        neighbor_valid_node = state.add_access(neighbor_valid_var, debuginfo=di)
+
+        neighbor_valid_tasklet = state.add_tasklet(
+            "check_valid_neighbor",
+            {"__idx"},
+            {"__valid"},
+            f"__valid = True if __idx != {neighbor_skip_value} else False",
+        )
+        state.add_edge(
+            neighbor_index_node,
+            None,
+            neighbor_valid_tasklet,
+            "__idx",
+            dace.Memlet(data=neighbor_index_var, subset="0"),
+        )
+        state.add_memlet_path(
+            neighbor_valid_tasklet,
+            mx,
+            neighbor_valid_node,
+            memlet=dace.Memlet(data=neighbor_valid_var, subset=neighbor_map_index),
+            src_conn="__valid",
+        )
+        return [
+            ValueExpr(neighbor_value_node, iterator.dtype),
+            ValueExpr(neighbor_valid_node, dace.dtypes.bool),
+        ]
 
 
 def builtin_can_deref(
@@ -419,6 +474,42 @@ def builtin_cast(
     )
 
 
+def builtin_make_const_list(
+    transformer: "PythonTaskletCodegen", node: itir.Expr, node_args: list[itir.Expr]
+) -> list[ValueExpr]:
+    di = dace_debuginfo(node, transformer.context.body.debuginfo)
+    args = [transformer.visit(arg)[0] for arg in node_args]
+    assert all(isinstance(x, (SymbolExpr, ValueExpr)) for x in args)
+    args_dtype = [x.dtype for x in args]
+    assert len(set(args_dtype)) == 1
+    dtype = args_dtype[0]
+
+    var_name = unique_var_name()
+    transformer.context.body.add_array(var_name, (len(args),), dtype, transient=True)
+    var_node = transformer.context.state.add_access(var_name, debuginfo=di)
+
+    for i, arg in enumerate(args):
+        if isinstance(arg, SymbolExpr):
+            transformer.context.state.add_edge(
+                transformer.context.state.add_tasklet(
+                    f"get_arg{i}", {}, {"val"}, f"val = {arg.value}"
+                ),
+                "val",
+                var_node,
+                None,
+                dace.Memlet(data=var_name, subset=f"{i}"),
+            )
+        else:
+            assert arg.value.desc(transformer.context.body).shape == (1,)
+            transformer.context.state.add_nedge(
+                arg.value,
+                var_node,
+                dace.Memlet(data=arg.value.data, subset="0", other_subset=f"{i}"),
+            )
+
+    return [ValueExpr(var_node, dtype)]
+
+
 def builtin_make_tuple(
     transformer: "PythonTaskletCodegen", node: itir.Expr, node_args: list[itir.Expr]
 ) -> list[ValueExpr]:
@@ -443,6 +534,7 @@ _GENERAL_BUILTIN_MAPPING: dict[
     "cast_": builtin_cast,
     "if_": builtin_if,
     "list_get": builtin_list_get,
+    "make_const_list": builtin_make_const_list,
     "make_tuple": builtin_make_tuple,
     "neighbors": builtin_neighbors,
     "tuple_get": builtin_tuple_get,
@@ -578,7 +670,7 @@ class PythonTaskletCodegen(gt4py.eve.codegen.TemplatedGenerator):
 
         # Create the SDFG for the lambda's body
         lambda_sdfg = dace.SDFG(func_name)
-        lambda_sdfg.debuginfo = dace_debuginfo(node)
+        lambda_sdfg.debuginfo = dace_debuginfo(node, self.context.body.debuginfo)
         lambda_state = lambda_sdfg.add_state(f"{func_name}_entry", True)
 
         lambda_symbols_pass = GatherLambdaSymbolsPass(
@@ -951,7 +1043,7 @@ class PythonTaskletCodegen(gt4py.eve.codegen.TemplatedGenerator):
 
             args = self.visit(node.args[0])
 
-            assert len(args) == 1
+            assert 1 <= len(args) <= 2
             reduce_input_node = args[0].value
 
         else:
@@ -967,16 +1059,22 @@ class PythonTaskletCodegen(gt4py.eve.codegen.TemplatedGenerator):
             # set reduction state in visit context
             self.context.reduce_identity = SymbolExpr(reduce_identity, reduce_dtype)
 
-            args = flatten_list(self.visit(node.args))
+            args = self.visit(node.args)
 
             # clear context
             self.context.reduce_identity = None
 
             # check that all neighbor expressions have the same shape
-            nreduce_shape = args[1].value.desc(self.context.body).shape
-            assert all(
-                [arg.value.desc(self.context.body).shape == nreduce_shape for arg in args[2:]]
-            )
+            args_shape = [
+                arg[0].value.desc(self.context.body).shape
+                for arg in args
+                if arg[0].value.desc(self.context.body).shape != (1,)
+            ]
+            assert len(set(args_shape)) == 1
+            nreduce_shape = args_shape[0]
+
+            input_args = [arg[0] for arg in args]
+            input_valid = [arg[1] for arg in args if len(arg) == 2]
 
             nreduce_index = tuple(f"_i{i}" for i in range(len(nreduce_shape)))
             nreduce_domain = {idx: f"0:{size}" for idx, size in zip(nreduce_index, nreduce_shape)}
@@ -990,12 +1088,16 @@ class PythonTaskletCodegen(gt4py.eve.codegen.TemplatedGenerator):
                 expr=fun_node.expr.args[1], params=fun_node.params[1:], location=node.location
             )
             lambda_context, inner_inputs, inner_outputs = self.visit(
-                lambda_node, args=args, use_neighbor_tables=False
+                lambda_node, args=input_args, use_neighbor_tables=False
             )
 
             input_mapping = {
-                param: create_memlet_at(arg.value.data, nreduce_index)
-                for (param, _), arg in zip(inner_inputs, args)
+                param: (
+                    dace.Memlet(data=arg.value.data, subset="0")
+                    if arg.value.desc(self.context.body).shape == (1,)
+                    else create_memlet_at(arg.value.data, nreduce_index)
+                )
+                for (param, _), arg in zip(inner_inputs, input_args)
             }
             output_mapping = {
                 inner_outputs[0].value.data: create_memlet_at(reduce_input_name, nreduce_index)
@@ -1004,6 +1106,42 @@ class PythonTaskletCodegen(gt4py.eve.codegen.TemplatedGenerator):
                 self.context.body, lambda_context.body, input_mapping
             )
 
+            if input_valid:
+                """
+                The neighbors builtin returns an array of booleans in case the connectivity table
+                contains skip values. These boolean values indicate whether the neighbor value is present or not,
+                and are used below to construct an if/else branch to bypass the lambda call for neighbor skip values.
+                If the neighbor table has full connectivity (no skip values by type definition), the input_valid node
+                is not built, and the construction of the if/else branch below is also skipped.
+                """
+                input_args.append(input_valid[0])
+                input_valid_node = input_valid[0].value
+                # add input connector to nested sdfg
+                input_mapping["is_valid"] = create_memlet_at(input_valid_node.data, nreduce_index)
+                # check neighbor validity on if/else inter-state edge
+                start_state = lambda_context.body.add_state("start", is_start_block=True)
+                skip_neighbor_state = lambda_context.body.add_state("skip_neighbor")
+                skip_neighbor_state.add_edge(
+                    skip_neighbor_state.add_tasklet(
+                        "identity", {}, {"val"}, f"val = {reduce_identity}"
+                    ),
+                    "val",
+                    skip_neighbor_state.add_access(inner_outputs[0].value.data),
+                    None,
+                    dace.Memlet(data=inner_outputs[0].value.data, subset="0"),
+                )
+                lambda_context.body.add_scalar("is_valid", dace.dtypes.bool)
+                lambda_context.body.add_edge(
+                    start_state,
+                    skip_neighbor_state,
+                    dace.InterstateEdge(condition="is_valid == False"),
+                )
+                lambda_context.body.add_edge(
+                    start_state,
+                    lambda_context.state,
+                    dace.InterstateEdge(condition="is_valid == True"),
+                )
+
             reduce_input_node = self.context.state.add_access(reduce_input_name, debuginfo=di)
 
             nsdfg_node, map_entry, _ = add_mapped_nested_sdfg(
@@ -1013,7 +1151,7 @@ class PythonTaskletCodegen(gt4py.eve.codegen.TemplatedGenerator):
                 inputs=input_mapping,
                 outputs=output_mapping,
                 symbol_mapping=symbol_mapping,
-                input_nodes={arg.value.data: arg.value for arg in args},
+                input_nodes={arg.value.data: arg.value for arg in input_args},
                 output_nodes={reduce_input_name: reduce_input_node},
                 debuginfo=di,
             )
@@ -1036,10 +1174,6 @@ class PythonTaskletCodegen(gt4py.eve.codegen.TemplatedGenerator):
             reduce_node, result_access, dace.Memlet(data=result_name, subset="0")
         )
 
-        # we apply map fusion only to the nested-SDFG which is generated for the reduction operator
-        # the purpose is to keep the ITIR-visitor program simple and to clean up the generated SDFG
-        self.context.body.apply_transformations_repeated([MapFusion], validate=False)
-
         return [ValueExpr(result_access, reduce_dtype)]
 
     def _visit_numeric_builtin(self, node: itir.FunCall) -> list[ValueExpr]:


Test patch:
diff --git a/tests/next_tests/definitions.py b/tests/next_tests/definitions.py
index c95292d7..56b220e0 100644
--- a/tests/next_tests/definitions.py
+++ b/tests/next_tests/definitions.py
@@ -155,7 +155,6 @@ COMMON_SKIP_TEST_LIST = [
     (USES_SPARSE_FIELDS_AS_OUTPUT, XFAIL, UNSUPPORTED_MESSAGE),
 ]
 DACE_SKIP_TEST_LIST = COMMON_SKIP_TEST_LIST + [
-    (USES_CONSTANT_FIELDS, XFAIL, UNSUPPORTED_MESSAGE),
     (USES_INDEX_FIELDS, XFAIL, UNSUPPORTED_MESSAGE),
     (USES_LIFT_EXPRESSIONS, XFAIL, UNSUPPORTED_MESSAGE),
     (USES_ORIGIN, XFAIL, UNSUPPORTED_MESSAGE),
diff --git a/tests/next_tests/integration_tests/feature_tests/iterator_tests/test_tuple.py b/tests/next_tests/integration_tests/feature_tests/iterator_tests/test_tuple.py
index add772e7..925ad33e 100644
--- a/tests/next_tests/integration_tests/feature_tests/iterator_tests/test_tuple.py
+++ b/tests/next_tests/integration_tests/feature_tests/iterator_tests/test_tuple.py
@@ -311,7 +311,9 @@ def test_tuple_field_input(program_processor):
     )
     inp2 = gtx.as_field(
         [IDim, JDim, KDim],
-        rng.normal(size=(shape[0], shape[1], shape[2])),
+        rng.normal(
+            size=(shape[0], shape[1], shape[2] + 1)
+        ),  # TODO(havogt) currently we allow different sizes, needed for icon4py compatibility
     )
 
     out = gtx.as_field([IDim, JDim, KDim], np.zeros(shape))
@@ -323,7 +325,7 @@ def test_tuple_field_input(program_processor):
     }
     run_processor(tuple_input[dom], program_processor, (inp1, inp2), out=out, offset_provider={})
     if validate:
-        assert np.allclose(inp1.asnumpy() + inp2.asnumpy(), out.asnumpy())
+        assert np.allclose(inp1.asnumpy() + inp2.asnumpy()[:, :, :-1], out.asnumpy())
 
 
 @pytest.mark.xfail(reason="Implement wrapper for extradim as tuple")