Skip to content

Commit

Permalink
[bug] MatrixType bug fix: Add additional restrictions for unpacking a…
Browse files Browse the repository at this point in the history
… Matrix (#6795)

Issue: #5819

### Brief Summary
  • Loading branch information
jim19930609 committed Dec 5, 2022
1 parent 61effe9 commit 354e7de
Show file tree
Hide file tree
Showing 2 changed files with 17 additions and 2 deletions.
4 changes: 4 additions & 0 deletions python/taichi/lang/ast/ast_transformer.py
Original file line number Diff line number Diff line change
Expand Up @@ -153,6 +153,10 @@ def build_assign_unpack(ctx, node_target, values, is_static_assign):

# Unpack: a, b, c = ti.Vector([1., 2., 3.])
if isinstance(values, impl.Expr) and values.ptr.is_tensor():
if len(values.get_shape()) > 1:
raise ValueError(
'Matrices with more than one columns cannot be unpacked')

values = ctx.ast_builder.expand_expr([values.ptr])
if len(values) == 1:
values = values[0]
Expand Down
15 changes: 13 additions & 2 deletions tests/python/test_tuple_assign.py
Original file line number Diff line number Diff line change
Expand Up @@ -207,8 +207,7 @@ def func():
func()


@test_utils.test(arch=get_host_arch_list())
def test_unpack_mismatch_matrix():
def _test_unpack_mismatch_matrix():
a = ti.field(ti.f32, ())
b = ti.field(ti.f32, ())
c = ti.field(ti.f32, ())
Expand All @@ -223,6 +222,18 @@ def func():
func()


@test_utils.test(arch=get_host_arch_list())
def test_unpack_mismatch_matrix():
_test_unpack_mismatch_matrix()


@test_utils.test(arch=get_host_arch_list(),
real_matrix=True,
real_matrix_scalarize=True)
def test_unpack_mismatch_matrix_scalarize():
_test_unpack_mismatch_matrix()


@test_utils.test(arch=get_host_arch_list())
def test_unpack_from_shape():
a = ti.field(ti.f32, ())
Expand Down

0 comments on commit 354e7de

Please sign in to comment.