| 412 | |
| 413 | |
| 414 | def test_symbolic_shape(): |
| 415 | @R.function |
| 416 | def foo(x: R.Tensor(("m", "n"), "float32")) -> R.Tensor(("m", "n"), "float32"): |
| 417 | m = T.int64() |
| 418 | n = T.int64() |
| 419 | gv0 = R.call_dps_packed("extern_func", x, R.Tensor((m, n), dtype="float32")) |
| 420 | return gv0 |
| 421 | |
| 422 | @R.function |
| 423 | def bar(x: R.Tensor(("m", "n"), "float32")) -> R.Tensor(("m", "n"), "float32"): |
| 424 | m = T.int64() |
| 425 | n = T.int64() |
| 426 | gv0 = R.call_dps_packed("extern_func", x, R.Tensor((m, n), dtype="float32")) |
| 427 | return gv0 |
| 428 | |
| 429 | with pytest.raises(tvm.error.DiagnosticError): |
| 430 | |
| 431 | @R.function |
| 432 | def mismatch_dtype(x: R.Tensor(("m", "n"), "float32")) -> R.Tensor(None, "float32", ndim=2): |
| 433 | m = T.int64() |
| 434 | n = T.int32() # The shape dtype should be int64 |
| 435 | gv0 = R.call_dps_packed("extern_func", x, R.Tensor((m, n), dtype="float32")) |
| 436 | return gv0 |
| 437 | |
| 438 | def _expected(name: str): |
| 439 | n, m = tirx.Var("n", "int64"), tirx.Var("m", "int64") |
| 440 | x = relax.Var("x", R.Tensor([m, n], "float32")) |
| 441 | bb = relax.BlockBuilder() |
| 442 | with bb.function(name, (x,)): |
| 443 | out = bb.emit( |
| 444 | relax.call_dps_packed("extern_func", x, R.Tensor((m, n), dtype="float32")) |
| 445 | ) |
| 446 | bb.emit_func_output(out) |
| 447 | return bb.get()[name] |
| 448 | |
| 449 | _check(foo, _expected("foo")) |
| 450 | _check(bar, _expected("bar")) |
| 451 | |
| 452 | |
| 453 | def test_shadowing(): |