Skip to content

Commit

Permalink
[dace] Update test case (1)
Browse files Browse the repository at this point in the history
  • Loading branch information
edopao committed Feb 26, 2024
1 parent 6e1a0c7 commit b8bcd46
Showing 1 changed file with 2 additions and 2 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -1067,7 +1067,7 @@ def test_where_k_offset(cartesian_case):
def fieldop_where_k_offset(
inp: cases.IKField, k_index: gtx.Field[[KDim], gtx.IndexType]
) -> cases.IKField:
return where(k_index > 0, inp(Koff[-1]), 2)
return where(k_index > 0, inp(Koff[-1]), -1)

@gtx.program
def prog(inp: cases.IKField, k_index: gtx.Field[[KDim], gtx.IndexType], out: cases.IKField):
Expand All @@ -1079,7 +1079,7 @@ def prog(inp: cases.IKField, k_index: gtx.Field[[KDim], gtx.IndexType], out: cas
)()
out = cases.allocate(cartesian_case, fieldop_where_k_offset, "inp")()

ref = np.where(k_index.asnumpy() > 0, np.roll(inp.asnumpy(), 1, axis=1), 2)
ref = np.where(k_index.asnumpy() > 0, np.roll(inp.asnumpy(), 1, axis=1), -1)

cases.verify(cartesian_case, prog, inp, k_index, out=out, ref=ref)

Expand Down

0 comments on commit b8bcd46

Please sign in to comment.