()
| 56 | |
| 57 | |
| 58 | def test_reshape_infer_struct_info(): |
| 59 | bb = relax.BlockBuilder() |
| 60 | vdev0 = VDevice("llvm") |
| 61 | x0 = relax.Var("x", R.Tensor((2, 3, 4, 5), "float32")) |
| 62 | x1 = relax.Var("x", R.Tensor("float32", ndim=4)) |
| 63 | x2 = relax.Var("x", R.Tensor("float32")) |
| 64 | x3 = relax.Var("x", R.Tensor((2, 3, 4, 5))) |
| 65 | x4 = relax.Var("x", R.Tensor(ndim=4)) |
| 66 | x5 = relax.Var("x", R.Tensor()) |
| 67 | x6 = relax.Var("x", R.Tensor((2, 3, 4, 5), "float32", vdev0)) |
| 68 | s0 = relax.Var("s", R.Shape((3, 8, 5))) |
| 69 | s1 = relax.Var("s", R.Shape(ndim=3)) |
| 70 | s2 = relax.Var("s", R.Shape()) |
| 71 | s3 = relax.ShapeExpr((3, 8, 5)) |
| 72 | |
| 73 | _check_inference( |
| 74 | bb, relax.op.reshape(x0, (3, 8, 5)), relax.TensorStructInfo((3, 8, 5), "float32") |
| 75 | ) |
| 76 | _check_inference( |
| 77 | bb, relax.op.reshape(x6, (3, 8, 5)), relax.TensorStructInfo((3, 8, 5), "float32", vdev0) |
| 78 | ) |
| 79 | _check_inference( |
| 80 | bb, relax.op.reshape(x0, (3, -1, 5)), relax.TensorStructInfo((3, 8, 5), "float32") |
| 81 | ) |
| 82 | _check_inference(bb, relax.op.reshape(x0, (-1,)), relax.TensorStructInfo((120,), "float32")) |
| 83 | _check_inference( |
| 84 | bb, relax.op.reshape(x0, relax.ShapeExpr([-1])), relax.TensorStructInfo((120,), "float32") |
| 85 | ) |
| 86 | _check_inference( |
| 87 | bb, relax.op.reshape(x1, (3, 8, 5)), relax.TensorStructInfo((3, 8, 5), "float32") |
| 88 | ) |
| 89 | _check_inference( |
| 90 | bb, relax.op.reshape(x2, (3, 8, 5)), relax.TensorStructInfo((3, 8, 5), "float32") |
| 91 | ) |
| 92 | _check_inference( |
| 93 | bb, relax.op.reshape(x3, (3, 8, 5)), relax.TensorStructInfo((3, 8, 5), dtype="") |
| 94 | ) |
| 95 | _check_inference( |
| 96 | bb, relax.op.reshape(x3, (3, -1, 5)), relax.TensorStructInfo((3, 8, 5), dtype="") |
| 97 | ) |
| 98 | _check_inference( |
| 99 | bb, relax.op.reshape(x4, (3, 8, 5)), relax.TensorStructInfo((3, 8, 5), dtype="") |
| 100 | ) |
| 101 | _check_inference( |
| 102 | bb, relax.op.reshape(x5, (3, 8, 5)), relax.TensorStructInfo((3, 8, 5), dtype="") |
| 103 | ) |
| 104 | # Remove Var from StructInfo when we can |
| 105 | _check_inference(bb, relax.op.reshape(x0, s0), relax.TensorStructInfo((3, 8, 5), "float32")) |
| 106 | _check_inference(bb, relax.op.reshape(x1, s0), relax.TensorStructInfo((3, 8, 5), "float32")) |
| 107 | _check_inference(bb, relax.op.reshape(x2, s0), relax.TensorStructInfo((3, 8, 5), "float32")) |
| 108 | _check_inference(bb, relax.op.reshape(x3, s0), relax.TensorStructInfo((3, 8, 5), dtype="")) |
| 109 | _check_inference(bb, relax.op.reshape(x4, s0), relax.TensorStructInfo((3, 8, 5), dtype="")) |
| 110 | _check_inference(bb, relax.op.reshape(x5, s0), relax.TensorStructInfo((3, 8, 5), dtype="")) |
| 111 | _check_inference(bb, relax.op.reshape(x0, s1), relax.TensorStructInfo(s1, "float32")) |
| 112 | _check_inference(bb, relax.op.reshape(x1, s1), relax.TensorStructInfo(s1, "float32")) |
| 113 | _check_inference(bb, relax.op.reshape(x2, s1), relax.TensorStructInfo(s1, "float32")) |
| 114 | _check_inference(bb, relax.op.reshape(x3, s1), relax.TensorStructInfo(s1, dtype="")) |
| 115 | _check_inference(bb, relax.op.reshape(x4, s1), relax.TensorStructInfo(s1, dtype="")) |
nothing calls this directly
no test coverage detected
searching dependent graphs…