| 227 | |
| 228 | |
| 229 | def test_shape_of(): |
| 230 | v0 = rx.Var("v0", R.Tensor(ndim=2)) |
| 231 | s0 = rx.get_shape_of(v0) |
| 232 | s0_str = dump_ast(s0) |
| 233 | assert s0_str.startswith("Call(") |
| 234 | assert 'op=Op(name="relax.shape_of")' in s0_str |
| 235 | assert "args=" in s0_str |
| 236 | assert 'name_hint="v0"' in s0_str |
| 237 | |
| 238 | v1 = rx.Var("v1", R.Tensor([96, 54])) |
| 239 | s1 = rx.get_shape_of(v1) |
| 240 | s1_str = dump_ast(s1) |
| 241 | assert s1_str.startswith("ShapeExpr("), s1_str |
| 242 | assert "values=" in s1_str |
| 243 | assert "PrimExpr(value=`T.int64(96)`)" in s1_str |
| 244 | assert "PrimExpr(value=`T.int64(54)`)" in s1_str |
| 245 | |
| 246 | |
| 247 | def test_shape_expr(): |