def test_vector_assign(left, is_serial: bool) -> None:
result = NpirCodegen().visit(
VectorAssignFactory(left=left, right=FieldSliceFactory(name="right")),
ctx=NpirCodegen.BlockContext(),
is_serial=is_serial,
)
left_str, right_str = result.split(" = ")
k_str = "k_:k_+1" if is_serial else "k:K"
if isinstance(left, npir.FieldSlice):
assert left_str == "left[i:I, j:J, " + k_str + "]"
else:
assert left_str == "left"
assert right_str == "right[i:I, j:J, " + k_str + "]"
def test_vector_assign(left, is_serial: bool) -> None:
result = NpirCodegen().visit(
VectorAssignFactory(left=left,
right=FieldSliceFactory(name="right",
k_offset=-1)),
ctx=NpirCodegen.BlockContext(),
is_serial=is_serial,
)
left_str, right_str = result.split(" = ")
if isinstance(left, npir.FieldSlice):
k_str_left = "k_:k_ + 1" if is_serial else "k:K"
else:
k_str_left = ":" if is_serial else "k:K"
k_str_right = "k_ - 1:k_" if is_serial else "k - 1:K - 1"
assert left_str == f"left[i:I, j:J, {k_str_left}]"
assert right_str == f"right[i:I, j:J, {k_str_right}]"