| 367 | |
| 368 | |
| 369 | def test_complex_seq_body(): |
| 370 | # Error: seq expr with a body that is not a leaf expression is not permitted |
| 371 | x = rx.Var("x", R.Tensor([], "int32")) |
| 372 | y = rx.Var("y", R.Tensor([], "int32")) |
| 373 | func = rx.Function( |
| 374 | [x, y], |
| 375 | rx.SeqExpr([], rx.op.add(x, y)), |
| 376 | R.Tensor(ndim=0, dtype="int32"), |
| 377 | ).with_attr("global_symbol", "foo") |
| 378 | mod = tvm.IRModule.from_expr(func) |
| 379 | assert not rx.analysis.check_well_formed(mod, check_struct_info=False) |
| 380 | |
| 381 | # but if the result is bound, then it's okay |
| 382 | z = rx.Var("z", R.Tensor([], "int32")) |
| 383 | new_func = rx.Function( |
| 384 | [x, y], |
| 385 | rx.SeqExpr( |
| 386 | [ |
| 387 | rx.BindingBlock( |
| 388 | [ |
| 389 | rx.VarBinding( |
| 390 | var=z, |
| 391 | value=rx.op.add(x, y), |
| 392 | ) |
| 393 | ] |
| 394 | ) |
| 395 | ], |
| 396 | z, |
| 397 | ), |
| 398 | R.Tensor(ndim=0, dtype="int32"), |
| 399 | ).with_attr("global_symbol", "foo") |
| 400 | new_mod = tvm.IRModule.from_expr(new_func) |
| 401 | # normalize in order to fill in checked type |
| 402 | normalized = rx.transform.Normalize()(new_mod) |
| 403 | rx.analysis.well_formed(normalized, check_struct_info=True) |
| 404 | |
| 405 | |
| 406 | def test_inline_prim_func(): |