wehub-resource-sync
3501 lines
134 KiB
Python
3501 lines
134 KiB
Python
# Licensed to the Apache Software Foundation (ASF) under one
|
|
# or more contributor license agreements. See the NOTICE file
|
|
# distributed with this work for additional information
|
|
# regarding copyright ownership. The ASF licenses this file
|
|
# to you under the Apache License, Version 2.0 (the
|
|
# "License"); you may not use this file except in compliance
|
|
# with the License. You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing,
|
|
# software distributed under the License is distributed on an
|
|
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
|
# KIND, either express or implied. See the License for the
|
|
# specific language governing permissions and limitations
|
|
# under the License.
|
|
# ruff: noqa: E731, F841
|
|
import pytest
|
|
|
|
import tvm
|
|
import tvm.testing
|
|
from tvm import relax, tirx
|
|
from tvm.ir import Op, VDevice
|
|
from tvm.script import relax as R
|
|
from tvm.script import tirx as T
|
|
|
|
|
|
def test_op_correctness():
|
|
x = relax.Var("x", R.Tensor((3, 4, 5), "float32"))
|
|
assert relax.op.broadcast_to(x, (3, 3, 4, 5)).op == Op.get("relax.broadcast_to")
|
|
assert relax.op.concat([x]).op == Op.get("relax.concat")
|
|
assert relax.op.expand_dims(x, axis=[]).op == Op.get("relax.expand_dims")
|
|
assert relax.op.flatten(x).op == Op.get("relax.flatten")
|
|
assert relax.op.permute_dims(x).op == Op.get("relax.permute_dims")
|
|
assert relax.op.reshape(x, (4, 5, 3)).op == Op.get("relax.reshape")
|
|
assert relax.op.split(x, indices_or_sections=1).op == Op.get("relax.split")
|
|
assert relax.op.tile(x, (2, 2, 2)).op == Op.get("relax.tile")
|
|
assert relax.op.repeat(x, 2, 0).op == Op.get("relax.repeat")
|
|
assert relax.op.squeeze(x).op == Op.get("relax.squeeze")
|
|
assert relax.op.layout_transform(x, index_map=lambda a, b, c: (b, c, a)).op == Op.get(
|
|
"relax.layout_transform"
|
|
)
|
|
assert relax.op.collapse_sum_to(x, (4, 5)).op == Op.get("relax.collapse_sum_to")
|
|
y = relax.Var("x", R.Tensor((4, 5), "float32"))
|
|
assert relax.op.collapse_sum_like(x, y).op == Op.get("relax.collapse_sum_like")
|
|
assert relax.op.cumsum(x, axis=1, dtype="int32").op == Op.get("relax.cumsum")
|
|
assert relax.op.einsum(x, subscripts="ii").op == Op.get("relax.einsum")
|
|
assert relax.op.flip(x, axis=1).op == Op.get("relax.flip")
|
|
seq_lengths = relax.Var("seq_lengths", R.Tensor((3,), "int32"))
|
|
assert relax.op.reverse_sequence(x, seq_lengths).op == Op.get("relax.reverse_sequence")
|
|
assert relax.op.scatter_elements(x, x, x).op == Op.get("relax.scatter_elements")
|
|
assert relax.op.scatter_nd(x, x, x).op == Op.get("relax.scatter_nd")
|
|
|
|
|
|
def _check_inference(bb: relax.BlockBuilder, call: relax.Call, expected_ty: relax.Type):
|
|
ret = bb.normalize(call)
|
|
tvm.ir.assert_structural_equal(ret.ty, expected_ty)
|
|
|
|
|
|
def test_reshape_infer_ty():
|
|
bb = relax.BlockBuilder()
|
|
vdev0 = VDevice("llvm")
|
|
x0 = relax.Var("x", R.Tensor((2, 3, 4, 5), "float32"))
|
|
x1 = relax.Var("x", R.Tensor("float32", ndim=4))
|
|
x2 = relax.Var("x", R.Tensor("float32"))
|
|
x3 = relax.Var("x", R.Tensor((2, 3, 4, 5)))
|
|
x4 = relax.Var("x", R.Tensor(ndim=4))
|
|
x5 = relax.Var("x", R.Tensor())
|
|
x6 = relax.Var("x", R.Tensor((2, 3, 4, 5), "float32", vdev0))
|
|
s0 = relax.Var("s", R.Shape((3, 8, 5)))
|
|
s1 = relax.Var("s", R.Shape(ndim=3))
|
|
s2 = relax.Var("s", R.Shape())
|
|
s3 = relax.ShapeExpr((3, 8, 5))
|
|
|
|
_check_inference(bb, relax.op.reshape(x0, (3, 8, 5)), relax.TensorType((3, 8, 5), "float32"))
|
|
_check_inference(
|
|
bb, relax.op.reshape(x6, (3, 8, 5)), relax.TensorType((3, 8, 5), "float32", vdev0)
|
|
)
|
|
_check_inference(bb, relax.op.reshape(x0, (3, -1, 5)), relax.TensorType((3, 8, 5), "float32"))
|
|
_check_inference(bb, relax.op.reshape(x0, (-1,)), relax.TensorType((120,), "float32"))
|
|
_check_inference(
|
|
bb, relax.op.reshape(x0, relax.ShapeExpr([-1])), relax.TensorType((120,), "float32")
|
|
)
|
|
_check_inference(bb, relax.op.reshape(x1, (3, 8, 5)), relax.TensorType((3, 8, 5), "float32"))
|
|
_check_inference(bb, relax.op.reshape(x2, (3, 8, 5)), relax.TensorType((3, 8, 5), "float32"))
|
|
_check_inference(bb, relax.op.reshape(x3, (3, 8, 5)), relax.TensorType((3, 8, 5), dtype=""))
|
|
_check_inference(bb, relax.op.reshape(x3, (3, -1, 5)), relax.TensorType((3, 8, 5), dtype=""))
|
|
_check_inference(bb, relax.op.reshape(x4, (3, 8, 5)), relax.TensorType((3, 8, 5), dtype=""))
|
|
_check_inference(bb, relax.op.reshape(x5, (3, 8, 5)), relax.TensorType((3, 8, 5), dtype=""))
|
|
# Remove Var from Type when we can
|
|
_check_inference(bb, relax.op.reshape(x0, s0), relax.TensorType((3, 8, 5), "float32"))
|
|
_check_inference(bb, relax.op.reshape(x1, s0), relax.TensorType((3, 8, 5), "float32"))
|
|
_check_inference(bb, relax.op.reshape(x2, s0), relax.TensorType((3, 8, 5), "float32"))
|
|
_check_inference(bb, relax.op.reshape(x3, s0), relax.TensorType((3, 8, 5), dtype=""))
|
|
_check_inference(bb, relax.op.reshape(x4, s0), relax.TensorType((3, 8, 5), dtype=""))
|
|
_check_inference(bb, relax.op.reshape(x5, s0), relax.TensorType((3, 8, 5), dtype=""))
|
|
_check_inference(bb, relax.op.reshape(x0, s1), relax.TensorType(s1, "float32"))
|
|
_check_inference(bb, relax.op.reshape(x1, s1), relax.TensorType(s1, "float32"))
|
|
_check_inference(bb, relax.op.reshape(x2, s1), relax.TensorType(s1, "float32"))
|
|
_check_inference(bb, relax.op.reshape(x3, s1), relax.TensorType(s1, dtype=""))
|
|
_check_inference(bb, relax.op.reshape(x4, s1), relax.TensorType(s1, dtype=""))
|
|
_check_inference(bb, relax.op.reshape(x5, s1), relax.TensorType(s1, dtype=""))
|
|
_check_inference(bb, relax.op.reshape(x0, s2), relax.TensorType(s2, "float32"))
|
|
_check_inference(bb, relax.op.reshape(x1, s2), relax.TensorType(s2, "float32"))
|
|
_check_inference(bb, relax.op.reshape(x2, s2), relax.TensorType(s2, "float32"))
|
|
_check_inference(bb, relax.op.reshape(x3, s2), relax.TensorType(s2, dtype=""))
|
|
_check_inference(bb, relax.op.reshape(x4, s2), relax.TensorType(s2, dtype=""))
|
|
_check_inference(bb, relax.op.reshape(x5, s2), relax.TensorType(s2, dtype=""))
|
|
_check_inference(bb, relax.op.reshape(x0, s3), relax.TensorType(s3, "float32"))
|
|
_check_inference(bb, relax.op.reshape(x1, s3), relax.TensorType(s3, "float32"))
|
|
_check_inference(bb, relax.op.reshape(x2, s3), relax.TensorType(s3, "float32"))
|
|
_check_inference(bb, relax.op.reshape(x3, s3), relax.TensorType(s3, dtype=""))
|
|
_check_inference(bb, relax.op.reshape(x4, s3), relax.TensorType(s3, dtype=""))
|
|
_check_inference(bb, relax.op.reshape(x5, s3), relax.TensorType(s3, dtype=""))
|
|
|
|
|
|
def test_reshape_infer_ty_shape_symbolic():
|
|
bb = relax.BlockBuilder()
|
|
a = tirx.Var("a", "int64")
|
|
b = tirx.Var("b", "int64")
|
|
c = tirx.Var("c", "int64")
|
|
d = tirx.Var("d", "int64")
|
|
x = relax.Var("x", R.Tensor((a, b, c, d), "float32"))
|
|
s0 = relax.Var("s", R.Shape((c, a, d, b)))
|
|
s1 = relax.Var("s", R.Shape())
|
|
s2 = relax.ShapeExpr((c, a, d, b))
|
|
|
|
_check_inference(
|
|
bb, relax.op.reshape(x, (c, a, d, b)), relax.TensorType((c, a, d, b), "float32")
|
|
)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.reshape(x, (d, c, b, -1)),
|
|
relax.TensorType((d, c, b, a), "float32"),
|
|
)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.reshape(x, (1, -1, 1)),
|
|
relax.TensorType((1, a * b * c * d, 1), "float32"),
|
|
)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.reshape(x, (2, -1, a)),
|
|
relax.TensorType((2, tirx.floordiv(b * c * d, 2), a), "float32"),
|
|
)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.reshape(x, (c, -1, d, b)),
|
|
relax.TensorType((c, a, d, b), "float32"),
|
|
)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.reshape(x, (c, a * d, b)),
|
|
relax.TensorType((c, a * d, b), "float32"),
|
|
)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.reshape(x, (c, a * b * d, -1)),
|
|
relax.TensorType((c, a * b * d, 1), "float32"),
|
|
)
|
|
# Remove Var from Type when we can
|
|
_check_inference(bb, relax.op.reshape(x, s0), relax.TensorType((c, a, d, b), "float32"))
|
|
_check_inference(bb, relax.op.reshape(x, s1), relax.TensorType(s1, "float32"))
|
|
_check_inference(bb, relax.op.reshape(x, s2), relax.TensorType(s2, "float32"))
|
|
|
|
|
|
def test_reshape_infer_ty_shape_var():
|
|
bb = relax.BlockBuilder()
|
|
s0 = relax.Var("s", relax.ShapeType((2, 3, 4, 5)))
|
|
s1 = relax.Var("s", relax.ShapeType(ndim=4))
|
|
s2 = relax.Var("s", relax.ShapeType())
|
|
x0 = relax.Var("x", relax.TensorType(s0, "float32"))
|
|
x1 = relax.Var("x", relax.TensorType(s1, "float32"))
|
|
x2 = relax.Var("x", relax.TensorType(s2, "float32"))
|
|
ns0 = relax.Var("ns", relax.ShapeType((3, 8, 5)))
|
|
ns1 = relax.Var("ns", relax.ShapeType())
|
|
|
|
_check_inference(bb, relax.op.reshape(x0, (3, 8, 5)), relax.TensorType((3, 8, 5), "float32"))
|
|
_check_inference(
|
|
bb, relax.op.reshape(x0, (2, 3, 0, 5)), relax.TensorType((2, 3, 4, 5), "float32")
|
|
)
|
|
_check_inference(
|
|
bb, relax.op.reshape(x0, (1, 3, 0, -1)), relax.TensorType((1, 3, 4, 10), "float32")
|
|
)
|
|
_check_inference(bb, relax.op.reshape(x0, (3, -1, 5)), relax.TensorType((3, 8, 5), "float32"))
|
|
# Remove Var from Type when we can
|
|
_check_inference(bb, relax.op.reshape(x0, ns0), relax.TensorType((3, 8, 5), "float32"))
|
|
_check_inference(bb, relax.op.reshape(x0, ns1), relax.TensorType(ns1, "float32"))
|
|
_check_inference(bb, relax.op.reshape(x1, (3, 8, 5)), relax.TensorType((3, 8, 5), "float32"))
|
|
# Remove Var from Type when we can
|
|
_check_inference(bb, relax.op.reshape(x1, ns0), relax.TensorType((3, 8, 5), "float32"))
|
|
_check_inference(bb, relax.op.reshape(x1, ns1), relax.TensorType(ns1, "float32"))
|
|
_check_inference(bb, relax.op.reshape(x2, (3, 8, 5)), relax.TensorType((3, 8, 5), "float32"))
|
|
# Remove Var from Type when we can
|
|
_check_inference(bb, relax.op.reshape(x2, ns0), relax.TensorType((3, 8, 5), "float32"))
|
|
_check_inference(bb, relax.op.reshape(x2, ns1), relax.TensorType(ns1, "float32"))
|
|
|
|
|
|
def test_reshape_infer_ty_more_input_dtype():
|
|
bb = relax.BlockBuilder()
|
|
x0 = relax.Var("x", R.Tensor((2, 3, 4, 5), "float16"))
|
|
x1 = relax.Var("x", R.Tensor((2, 3, 4, 5), "int8"))
|
|
|
|
_check_inference(bb, relax.op.reshape(x0, (120,)), relax.TensorType((120,), "float16"))
|
|
_check_inference(bb, relax.op.reshape(x1, (120,)), relax.TensorType((120,), "int8"))
|
|
|
|
|
|
def test_reshape_infer_ty_unequal_shape_prod():
|
|
bb = relax.BlockBuilder()
|
|
s = relax.Var("s", relax.ShapeType((2, 3, 4, 5)))
|
|
x0 = relax.Var("x", R.Tensor((2, 3, 4, 5), "float32"))
|
|
x1 = relax.Var("x", relax.TensorType(s, "float32"))
|
|
ns = relax.Var("ns", relax.ShapeType((4, 4, 1, 5)))
|
|
|
|
with pytest.raises(ValueError):
|
|
bb.normalize(relax.op.reshape(x0, (4, 4, 1, 5)))
|
|
with pytest.raises(ValueError):
|
|
bb.normalize(relax.op.reshape(x1, (4, 4, 1, 5)))
|
|
with pytest.raises(ValueError):
|
|
bb.normalize(relax.op.reshape(x0, (4, 4, -1, 5)))
|
|
with pytest.raises(ValueError):
|
|
bb.normalize(relax.op.reshape(x1, (4, 4, -1, 5)))
|
|
with pytest.raises(ValueError):
|
|
bb.normalize(relax.op.reshape(x0, ns))
|
|
with pytest.raises(ValueError):
|
|
bb.normalize(relax.op.reshape(x1, ns))
|
|
|
|
|
|
def test_reshape_infer_ty_inference_not_deducible():
|
|
bb = relax.BlockBuilder()
|
|
s0 = relax.Var("s", relax.ShapeType(ndim=4))
|
|
s1 = relax.Var("s", relax.ShapeType())
|
|
x0 = relax.Var("x", R.Tensor("float32", ndim=4))
|
|
x1 = relax.Var("x", R.Tensor("float32"))
|
|
x2 = relax.Var("x", relax.TensorType(s0, "float32"))
|
|
x3 = relax.Var("x", relax.TensorType(s1, "float32"))
|
|
|
|
with pytest.raises(tvm.error.InternalError):
|
|
bb.normalize(relax.op.reshape(x0, (2, 3, -1)))
|
|
with pytest.raises(tvm.error.InternalError):
|
|
bb.normalize(relax.op.reshape(x1, (2, 3, -1)))
|
|
with pytest.raises(tvm.error.InternalError):
|
|
bb.normalize(relax.op.reshape(x2, (2, 3, -1)))
|
|
with pytest.raises(tvm.error.InternalError):
|
|
bb.normalize(relax.op.reshape(x3, (2, 3, -1)))
|
|
|
|
|
|
def test_reshape_new_shape_not_tuple():
|
|
x = relax.Var("x", R.Tensor((2, 3, 4, 5), "float32"))
|
|
|
|
with pytest.raises(TypeError):
|
|
relax.op.reshape(x, 120)
|
|
|
|
|
|
def test_reshape_infer_ty_new_shape_not_integer():
|
|
bb = relax.BlockBuilder()
|
|
x = relax.Var("x", R.Tensor((2, 3, 4, 5), "float32"))
|
|
|
|
with pytest.raises(tvm.error.InternalError):
|
|
bb.normalize(relax.op.reshape(x, (2.0, 3, 4, 5)))
|
|
with pytest.raises(tvm.error.InternalError):
|
|
bb.normalize(relax.op.reshape(x, (2, 3, -1.0)))
|
|
with pytest.raises(tvm.error.InternalError):
|
|
bb.normalize(relax.op.reshape(x, (2, 3, 4.0, -1)))
|
|
|
|
|
|
def test_reshape_infer_ty_multiple_dim_inference():
|
|
bb = relax.BlockBuilder()
|
|
x = relax.Var("x", R.Tensor((2, 3, 4, 5), "float32"))
|
|
|
|
with pytest.raises(tvm.error.InternalError):
|
|
bb.normalize(relax.op.reshape(x, (2, -1, -1, 5)))
|
|
with pytest.raises(tvm.error.InternalError):
|
|
bb.normalize(relax.op.reshape(x, (-1, -1, -1, -1)))
|
|
|
|
|
|
def test_reshape_infer_ty_non_positive_new_shape():
|
|
bb = relax.BlockBuilder()
|
|
x = relax.Var("x", R.Tensor((2, 3, 4, 5), "float32"))
|
|
|
|
with pytest.raises(tvm.error.InternalError):
|
|
bb.normalize(relax.op.reshape(x, (-2, -3, -4, -5)))
|
|
|
|
|
|
def test_reshape_infer_ty_wrong_input_type():
|
|
bb = relax.BlockBuilder()
|
|
x0 = relax.Var("x", relax.ShapeType((2, 3, 4, 5)))
|
|
x1 = relax.Var("x", relax.FuncType([], R.Tensor((2, 3, 4, 5), "float32")))
|
|
x2 = relax.Var("x", R.Tensor((2, 3, 4, 5), "float32"))
|
|
ns = relax.Var("ns", relax.TensorType((120,), "float32"))
|
|
pv = relax.Var("pv", tvm.ir.PrimType("int64"))
|
|
|
|
with pytest.raises(TypeError):
|
|
bb.normalize(relax.op.reshape(x0, (2, 3, 4, 5)))
|
|
with pytest.raises(TypeError):
|
|
bb.normalize(relax.op.reshape(x1, (2, 3, 4, 5)))
|
|
with pytest.raises(TypeError):
|
|
bb.normalize(relax.op.reshape(x2, ns))
|
|
with pytest.raises(TypeError):
|
|
bb.normalize(relax.op.reshape(x2, [pv]))
|
|
|
|
|
|
def test_permute_dims_infer_ty():
|
|
bb = relax.BlockBuilder()
|
|
vdev0 = VDevice("llvm")
|
|
x0 = relax.Var("x", R.Tensor((1, 2, 3, 4), "float32"))
|
|
x1 = relax.Var("x", R.Tensor("float32", ndim=4))
|
|
x2 = relax.Var("x", R.Tensor("float32"))
|
|
x3 = relax.Var("x", R.Tensor((1, 2, 3, 4)))
|
|
x4 = relax.Var("x", R.Tensor(ndim=4))
|
|
x5 = relax.Var("x", R.Tensor())
|
|
x6 = relax.Var("x", R.Tensor((1,), "float32"))
|
|
x7 = relax.Var("x", R.Tensor((), "float32"))
|
|
x8 = relax.Var("x", R.Tensor((1, 2, 3, 4), "float32", vdev0))
|
|
|
|
_check_inference(
|
|
bb, relax.op.permute_dims(x0, [2, 3, 1, 0]), relax.TensorType((3, 4, 2, 1), "float32")
|
|
)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.permute_dims(x8, [2, 3, 1, 0]),
|
|
relax.TensorType((3, 4, 2, 1), "float32", vdev0),
|
|
)
|
|
_check_inference(
|
|
bb, relax.op.permute_dims(x0, axes=None), relax.TensorType((4, 3, 2, 1), "float32")
|
|
)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.permute_dims(x0, [-2, -3, 3, -4]),
|
|
relax.TensorType((3, 2, 4, 1), "float32"),
|
|
)
|
|
_check_inference(
|
|
bb, relax.op.permute_dims(x1, [2, 3, 1, 0]), relax.TensorType(dtype="float32", ndim=4)
|
|
)
|
|
_check_inference(
|
|
bb, relax.op.permute_dims(x1, axes=None), relax.TensorType(dtype="float32", ndim=4)
|
|
)
|
|
_check_inference(bb, relax.op.permute_dims(x2, axes=None), relax.TensorType(dtype="float32"))
|
|
_check_inference(
|
|
bb, relax.op.permute_dims(x3, [2, 3, 1, 0]), relax.TensorType((3, 4, 2, 1), dtype="")
|
|
)
|
|
_check_inference(
|
|
bb, relax.op.permute_dims(x3, axes=None), relax.TensorType((4, 3, 2, 1), dtype="")
|
|
)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.permute_dims(x3, [-2, -3, 3, -4]),
|
|
relax.TensorType((3, 2, 4, 1), dtype=""),
|
|
)
|
|
_check_inference(
|
|
bb, relax.op.permute_dims(x4, [2, 3, 1, 0]), relax.TensorType(dtype="", ndim=4)
|
|
)
|
|
_check_inference(bb, relax.op.permute_dims(x4, axes=None), relax.TensorType(dtype="", ndim=4))
|
|
_check_inference(bb, relax.op.permute_dims(x5, axes=None), relax.TensorType(dtype=""))
|
|
_check_inference(bb, relax.op.permute_dims(x6, axes=None), relax.TensorType((1,), "float32"))
|
|
_check_inference(bb, relax.op.permute_dims(x7, axes=None), relax.TensorType((), "float32"))
|
|
|
|
|
|
def test_permute_dims_infer_ty_shape_symbolic():
|
|
bb = relax.BlockBuilder()
|
|
a = tirx.Var("a", "int64")
|
|
b = tirx.Var("b", "int64")
|
|
c = tirx.Var("c", "int64")
|
|
d = tirx.Var("d", "int64")
|
|
x = relax.Var("x", R.Tensor((a, b, c, d), "float32"))
|
|
|
|
_check_inference(
|
|
bb, relax.op.permute_dims(x, [2, 3, 1, 0]), relax.TensorType((c, d, b, a), "float32")
|
|
)
|
|
_check_inference(
|
|
bb, relax.op.permute_dims(x, axes=None), relax.TensorType((d, c, b, a), "float32")
|
|
)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.permute_dims(x, [-2, -3, 3, -4]),
|
|
relax.TensorType((c, b, d, a), "float32"),
|
|
)
|
|
|
|
|
|
def test_permute_dims_infer_ty_shape_var():
|
|
bb = relax.BlockBuilder()
|
|
s0 = relax.Var("s", relax.ShapeType((1, 2, 3, 4)))
|
|
s1 = relax.Var("s", relax.ShapeType(ndim=4))
|
|
s2 = relax.Var("s", relax.ShapeType())
|
|
x0 = relax.Var("x", relax.TensorType(s0, "float32"))
|
|
x1 = relax.Var("x", relax.TensorType(s1, "float32"))
|
|
x2 = relax.Var("x", relax.TensorType(s2, "float32"))
|
|
|
|
_check_inference(bb, relax.op.permute_dims(x0, [0, 1, 2, 3]), relax.TensorType(s0, "float32"))
|
|
_check_inference(
|
|
bb, relax.op.permute_dims(x0, [-4, -3, -2, -1]), relax.TensorType(s0, "float32")
|
|
)
|
|
_check_inference(
|
|
bb, relax.op.permute_dims(x0, [2, 3, 0, 1]), relax.TensorType(dtype="float32", ndim=4)
|
|
)
|
|
_check_inference(
|
|
bb, relax.op.permute_dims(x0, axes=None), relax.TensorType(dtype="float32", ndim=4)
|
|
)
|
|
_check_inference(bb, relax.op.permute_dims(x1, [0, 1, 2, 3]), relax.TensorType(s1, "float32"))
|
|
_check_inference(
|
|
bb, relax.op.permute_dims(x1, [2, 3, 0, 1]), relax.TensorType(dtype="float32", ndim=4)
|
|
)
|
|
_check_inference(
|
|
bb, relax.op.permute_dims(x1, axes=None), relax.TensorType(dtype="float32", ndim=4)
|
|
)
|
|
_check_inference(bb, relax.op.permute_dims(x2, axes=None), relax.TensorType(dtype="float32"))
|
|
|
|
|
|
def test_permute_dims_infer_ty_more_input_dtype():
|
|
bb = relax.BlockBuilder()
|
|
x0 = relax.Var("x", R.Tensor((1, 2, 3, 4), "float16"))
|
|
x1 = relax.Var("x", R.Tensor((1, 2, 3, 4), "int8"))
|
|
x2 = relax.Var("x", R.Tensor((1, 2, 3, 4), "int32"))
|
|
|
|
_check_inference(
|
|
bb, relax.op.permute_dims(x0, [2, 3, 1, 0]), relax.TensorType((3, 4, 2, 1), "float16")
|
|
)
|
|
_check_inference(
|
|
bb, relax.op.permute_dims(x1, [2, 3, 1, 0]), relax.TensorType((3, 4, 2, 1), "int8")
|
|
)
|
|
_check_inference(
|
|
bb, relax.op.permute_dims(x2, [2, 3, 1, 0]), relax.TensorType((3, 4, 2, 1), "int32")
|
|
)
|
|
|
|
|
|
def test_permute_dims_infer_ty_unknown_ndim_with_axes():
|
|
bb = relax.BlockBuilder()
|
|
s = relax.Var("s", relax.ShapeType())
|
|
x0 = relax.Var("x", R.Tensor("float32"))
|
|
x1 = relax.Var("x", relax.TensorType(s, "float32"))
|
|
|
|
with pytest.raises(tvm.error.InternalError):
|
|
bb.normalize(relax.op.permute_dims(x0, [2, 3, 1, 0]))
|
|
with pytest.raises(tvm.error.InternalError):
|
|
bb.normalize(relax.op.permute_dims(x1, [2, 3, 1, 0]))
|
|
|
|
|
|
def test_permute_dims_infer_ty_wrong_number_axes():
|
|
bb = relax.BlockBuilder()
|
|
s0 = relax.Var("s", relax.ShapeType((1, 2, 3, 4)))
|
|
s1 = relax.Var("s", relax.ShapeType(ndim=4))
|
|
x0 = relax.Var("x", R.Tensor((1, 2, 3, 4), "float32"))
|
|
x1 = relax.Var("x", R.Tensor("float32", ndim=4))
|
|
x2 = relax.Var("x", relax.TensorType(s0, "float32"))
|
|
x3 = relax.Var("x", relax.TensorType(s1, "float32"))
|
|
|
|
with pytest.raises(ValueError):
|
|
bb.normalize(relax.op.permute_dims(x0, [0, 2, 1]))
|
|
with pytest.raises(ValueError):
|
|
bb.normalize(relax.op.permute_dims(x0, [1, 2, 4, 0, 3]))
|
|
with pytest.raises(ValueError):
|
|
bb.normalize(relax.op.permute_dims(x1, [0, 2, 1]))
|
|
with pytest.raises(ValueError):
|
|
bb.normalize(relax.op.permute_dims(x1, [1, 2, 4, 0, 3]))
|
|
with pytest.raises(ValueError):
|
|
bb.normalize(relax.op.permute_dims(x2, [0, 2, 1]))
|
|
with pytest.raises(ValueError):
|
|
bb.normalize(relax.op.permute_dims(x2, [1, 2, 4, 0, 3]))
|
|
with pytest.raises(ValueError):
|
|
bb.normalize(relax.op.permute_dims(x3, [0, 2, 1]))
|
|
with pytest.raises(ValueError):
|
|
bb.normalize(relax.op.permute_dims(x3, [1, 2, 4, 0, 3]))
|
|
|
|
|
|
def test_permute_dims_infer_ty_axis_out_of_range():
|
|
bb = relax.BlockBuilder()
|
|
x0 = relax.Var("x", R.Tensor((1, 2, 3, 4), "float32"))
|
|
x1 = relax.Var("x", R.Tensor("float32", ndim=4))
|
|
|
|
with pytest.raises(ValueError):
|
|
bb.normalize(relax.op.permute_dims(x0, [0, 3, 4, 1]))
|
|
with pytest.raises(ValueError):
|
|
bb.normalize(relax.op.permute_dims(x0, [0, -5, 1, 3]))
|
|
with pytest.raises(ValueError):
|
|
bb.normalize(relax.op.permute_dims(x1, [0, 3, 4, 1]))
|
|
with pytest.raises(ValueError):
|
|
bb.normalize(relax.op.permute_dims(x1, [0, -5, 1, 3]))
|
|
|
|
|
|
def test_permute_dims_infer_ty_repetitive_axes():
|
|
bb = relax.BlockBuilder()
|
|
x0 = relax.Var("x", R.Tensor((1, 2, 3, 4), "float32"))
|
|
x1 = relax.Var("x", R.Tensor("float32", ndim=4))
|
|
|
|
with pytest.raises(ValueError):
|
|
bb.normalize(relax.op.permute_dims(x0, [0, 2, 2, 1]))
|
|
with pytest.raises(ValueError):
|
|
bb.normalize(relax.op.permute_dims(x0, [0, 2, -2, 1]))
|
|
with pytest.raises(ValueError):
|
|
bb.normalize(relax.op.permute_dims(x1, [0, 2, 2, 1]))
|
|
with pytest.raises(ValueError):
|
|
bb.normalize(relax.op.permute_dims(x1, [0, 2, -2, 1]))
|
|
|
|
|
|
def test_permute_dims_infer_ty_wrong_input_type():
|
|
bb = relax.BlockBuilder()
|
|
x0 = relax.Var("x", relax.ShapeType((1, 2, 3, 4)))
|
|
x1 = relax.Var("x", relax.FuncType([], R.Tensor((1, 2, 3, 4), "float32")))
|
|
|
|
with pytest.raises(TypeError):
|
|
bb.normalize(relax.op.permute_dims(x0))
|
|
with pytest.raises(TypeError):
|
|
bb.normalize(relax.op.permute_dims(x1))
|
|
|
|
|
|
def test_expand_dims_infer_ty():
|
|
bb = relax.BlockBuilder()
|
|
vdev0 = VDevice("llvm")
|
|
x0 = relax.Var("x", R.Tensor((2, 3, 4), "float32"))
|
|
x1 = relax.Var("x", R.Tensor("float32", ndim=3))
|
|
x2 = relax.Var("x", R.Tensor("float32"))
|
|
x3 = relax.Var("x", R.Tensor((2, 3, 4)))
|
|
x4 = relax.Var("x", R.Tensor(ndim=3))
|
|
x5 = relax.Var("x", R.Tensor())
|
|
x6 = relax.Var("x", R.Tensor((2, 3, 4), "float32", vdev0))
|
|
|
|
_check_inference(
|
|
bb, relax.op.expand_dims(x0, [1, 3]), relax.TensorType((2, 1, 3, 1, 4), "float32")
|
|
)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.expand_dims(x6, [1, 3]),
|
|
relax.TensorType((2, 1, 3, 1, 4), "float32", vdev0),
|
|
)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.expand_dims(x0, [-1, 1, -6, 3, 5]),
|
|
relax.TensorType((2, 1, 1, 1, 3, 1, 4, 1), "float32"),
|
|
)
|
|
_check_inference(bb, relax.op.expand_dims(x0, []), relax.TensorType((2, 3, 4), "float32"))
|
|
_check_inference(
|
|
bb, relax.op.expand_dims(x1, [1, 3]), relax.TensorType(dtype="float32", ndim=5)
|
|
)
|
|
_check_inference(bb, relax.op.expand_dims(x1, []), relax.TensorType(dtype="float32", ndim=3))
|
|
_check_inference(bb, relax.op.expand_dims(x2, [1, 3]), relax.TensorType(dtype="float32"))
|
|
_check_inference(bb, relax.op.expand_dims(x2, []), relax.TensorType(dtype="float32"))
|
|
_check_inference(
|
|
bb, relax.op.expand_dims(x3, [1, 3]), relax.TensorType((2, 1, 3, 1, 4), dtype="")
|
|
)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.expand_dims(x3, [-1, 1, -6, 3, 5]),
|
|
relax.TensorType((2, 1, 1, 1, 3, 1, 4, 1), dtype=""),
|
|
)
|
|
_check_inference(bb, relax.op.expand_dims(x3, []), relax.TensorType((2, 3, 4), dtype=""))
|
|
_check_inference(bb, relax.op.expand_dims(x4, [1, 3]), relax.TensorType(dtype="", ndim=5))
|
|
_check_inference(bb, relax.op.expand_dims(x4, []), relax.TensorType(dtype="", ndim=3))
|
|
_check_inference(bb, relax.op.expand_dims(x5, [1, 3]), relax.TensorType(dtype=""))
|
|
_check_inference(bb, relax.op.expand_dims(x5, []), relax.TensorType(dtype=""))
|
|
|
|
|
|
def test_expand_dims_infer_ty_shape_symbolic():
|
|
bb = relax.BlockBuilder()
|
|
a = tirx.Var("a", "int64")
|
|
b = tirx.Var("b", "int64")
|
|
x = relax.Var("x", R.Tensor((a, 4, b), "float32"))
|
|
|
|
_check_inference(
|
|
bb, relax.op.expand_dims(x, [1, 3]), relax.TensorType((a, 1, 4, 1, b), "float32")
|
|
)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.expand_dims(x, [-1, 1, -6, 3, 5]),
|
|
relax.TensorType((a, 1, 1, 1, 4, 1, b, 1), "float32"),
|
|
)
|
|
_check_inference(bb, relax.op.expand_dims(x, []), relax.TensorType((a, 4, b), "float32"))
|
|
|
|
|
|
def test_expand_dims_infer_ty_shape_var():
|
|
bb = relax.BlockBuilder()
|
|
s0 = relax.Var("s", relax.ShapeType((2, 3, 4)))
|
|
s1 = relax.Var("s", relax.ShapeType(ndim=3))
|
|
s2 = relax.Var("s", relax.ShapeType())
|
|
x0 = relax.Var("x", relax.TensorType(s0, "float32"))
|
|
x1 = relax.Var("x", relax.TensorType(s1, "float32"))
|
|
x2 = relax.Var("x", relax.TensorType(s2, "float32"))
|
|
|
|
_check_inference(
|
|
bb, relax.op.expand_dims(x0, [1, 3]), relax.TensorType(dtype="float32", ndim=5)
|
|
)
|
|
_check_inference(bb, relax.op.expand_dims(x0, []), relax.TensorType(s0, "float32"))
|
|
_check_inference(
|
|
bb, relax.op.expand_dims(x1, [1, 3]), relax.TensorType(dtype="float32", ndim=5)
|
|
)
|
|
_check_inference(bb, relax.op.expand_dims(x1, []), relax.TensorType(s1, "float32"))
|
|
_check_inference(bb, relax.op.expand_dims(x2, [1, 3]), relax.TensorType(dtype="float32"))
|
|
_check_inference(bb, relax.op.expand_dims(x2, []), relax.TensorType(s2, "float32"))
|
|
|
|
|
|
def test_expand_dims_infer_ty_more_input_dtype():
|
|
bb = relax.BlockBuilder()
|
|
x0 = relax.Var("x", R.Tensor((2, 3, 4), "float16"))
|
|
x1 = relax.Var("x", R.Tensor((2, 3, 4), "int8"))
|
|
x2 = relax.Var("x", R.Tensor((2, 3, 4), "int32"))
|
|
|
|
_check_inference(
|
|
bb, relax.op.expand_dims(x0, [1, 3]), relax.TensorType((2, 1, 3, 1, 4), "float16")
|
|
)
|
|
_check_inference(
|
|
bb, relax.op.expand_dims(x1, [1, 3]), relax.TensorType((2, 1, 3, 1, 4), "int8")
|
|
)
|
|
_check_inference(
|
|
bb, relax.op.expand_dims(x2, [1, 3]), relax.TensorType((2, 1, 3, 1, 4), "int32")
|
|
)
|
|
|
|
|
|
def test_expand_dims_infer_ty_axis_out_of_range():
|
|
bb = relax.BlockBuilder()
|
|
s0 = relax.Var("s", relax.ShapeType((2, 3, 4)))
|
|
s1 = relax.Var("s", relax.ShapeType(ndim=3))
|
|
x0 = relax.Var("x", R.Tensor((2, 3, 4), "float32"))
|
|
x1 = relax.Var("x", R.Tensor("float32", ndim=3))
|
|
x2 = relax.Var("x", relax.TensorType(s0))
|
|
x3 = relax.Var("x", relax.TensorType(s1))
|
|
|
|
with pytest.raises(ValueError):
|
|
bb.normalize(relax.op.expand_dims(x0, [1, 5]))
|
|
with pytest.raises(ValueError):
|
|
bb.normalize(relax.op.expand_dims(x0, [-6, 1]))
|
|
with pytest.raises(ValueError):
|
|
bb.normalize(relax.op.expand_dims(x1, [1, 5]))
|
|
with pytest.raises(ValueError):
|
|
bb.normalize(relax.op.expand_dims(x1, [-6, 1]))
|
|
with pytest.raises(ValueError):
|
|
bb.normalize(relax.op.expand_dims(x2, [1, 5]))
|
|
with pytest.raises(ValueError):
|
|
bb.normalize(relax.op.expand_dims(x2, [-6, 1]))
|
|
with pytest.raises(ValueError):
|
|
bb.normalize(relax.op.expand_dims(x3, [1, 5]))
|
|
with pytest.raises(ValueError):
|
|
bb.normalize(relax.op.expand_dims(x3, [-6, 1]))
|
|
|
|
|
|
def test_expand_dims_infer_ty_repetitive_axes():
|
|
bb = relax.BlockBuilder()
|
|
s0 = relax.Var("s", relax.ShapeType((2, 3, 4)))
|
|
s1 = relax.Var("s", relax.ShapeType(ndim=3))
|
|
x0 = relax.Var("x", R.Tensor((2, 3, 4), "float32"))
|
|
x1 = relax.Var("x", R.Tensor("float32", ndim=3))
|
|
x2 = relax.Var("x", relax.TensorType(s0))
|
|
x3 = relax.Var("x", relax.TensorType(s1))
|
|
|
|
with pytest.raises(ValueError):
|
|
bb.normalize(relax.op.expand_dims(x0, [1, 1]))
|
|
with pytest.raises(ValueError):
|
|
bb.normalize(relax.op.expand_dims(x0, [1, -4]))
|
|
with pytest.raises(ValueError):
|
|
bb.normalize(relax.op.expand_dims(x1, [1, 1]))
|
|
with pytest.raises(ValueError):
|
|
bb.normalize(relax.op.expand_dims(x1, [1, -4]))
|
|
with pytest.raises(ValueError):
|
|
bb.normalize(relax.op.expand_dims(x2, [1, 1]))
|
|
with pytest.raises(ValueError):
|
|
bb.normalize(relax.op.expand_dims(x2, [1, -4]))
|
|
with pytest.raises(ValueError):
|
|
bb.normalize(relax.op.expand_dims(x3, [1, 1]))
|
|
with pytest.raises(ValueError):
|
|
bb.normalize(relax.op.expand_dims(x3, [1, -4]))
|
|
|
|
|
|
def test_expand_dims_infer_ty_wrong_input_type():
|
|
bb = relax.BlockBuilder()
|
|
x0 = relax.Var("x", relax.ShapeType((2, 3, 4)))
|
|
x1 = relax.Var("x", relax.FuncType([], R.Tensor((2, 3, 4), "float32")))
|
|
|
|
with pytest.raises(TypeError):
|
|
bb.normalize(relax.op.expand_dims(x0, axis=[]))
|
|
with pytest.raises(TypeError):
|
|
bb.normalize(relax.op.expand_dims(x1, axis=[]))
|
|
|
|
|
|
def test_layout_transform_infer_ty():
|
|
bb = relax.BlockBuilder()
|
|
vdev0 = VDevice("llvm")
|
|
x = relax.Var("x", R.Tensor((10, 20, 30), "float32"))
|
|
x1 = relax.Var("x", R.Tensor((10, 20, 30), "float32", vdev0))
|
|
|
|
transpose_transform = lambda a, b, c: (a, c, b)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.layout_transform(x, index_map=transpose_transform),
|
|
relax.TensorType((10, 30, 20), "float32"),
|
|
)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.layout_transform(x1, index_map=transpose_transform),
|
|
relax.TensorType((10, 30, 20), "float32", vdev0),
|
|
)
|
|
|
|
tiling_transform = lambda a, b, c: (a, b // 2, c, b % 2)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.layout_transform(x, index_map=tiling_transform),
|
|
relax.TensorType((10, 10, 30, 2), "float32"),
|
|
)
|
|
|
|
implicit_padding_transform = lambda a, b, c: (a, c, b // 3, b % 3)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.layout_transform(x, index_map=implicit_padding_transform, pad_value=2),
|
|
relax.TensorType((10, 30, 7, 3), "float32"),
|
|
)
|
|
|
|
flatten_transform = lambda a, b, c: a * 600 + b * 30 + c
|
|
_check_inference(
|
|
bb,
|
|
relax.op.layout_transform(x, index_map=flatten_transform),
|
|
relax.TensorType((6000,), "float32"),
|
|
)
|
|
|
|
|
|
def test_layout_transform_infer_ty_mismatch_dtype():
|
|
bb = relax.BlockBuilder()
|
|
x = relax.Var("x", R.Tensor((10, 20, 30), "int32"))
|
|
|
|
transpose_transform = lambda a, b, c: (a, c, b)
|
|
with pytest.raises(TypeError):
|
|
bb.normalize(relax.op.layout_transform(x, index_map=transpose_transform, pad_value=2.2))
|
|
|
|
|
|
def test_layout_transform_infer_ty_unknown_shape():
|
|
bb = relax.BlockBuilder()
|
|
tiling_transform = lambda a, b: (a, b // 2, b % 2)
|
|
|
|
x_unknown_shape = relax.Var("x", R.Tensor("float32", ndim=2))
|
|
_check_inference(
|
|
bb,
|
|
relax.op.layout_transform(x_unknown_shape, index_map=tiling_transform),
|
|
relax.TensorType(dtype="float32", ndim=3),
|
|
)
|
|
|
|
x_unknown_rank_dtype = relax.Var("x", R.Tensor())
|
|
_check_inference(
|
|
bb,
|
|
relax.op.layout_transform(x_unknown_rank_dtype, index_map=tiling_transform),
|
|
relax.TensorType(dtype="", ndim=3),
|
|
)
|
|
|
|
|
|
def test_layout_transform_infer_ty_symbolic_shape():
|
|
bb = relax.BlockBuilder()
|
|
a = tirx.Var("a", "int64")
|
|
b = tirx.Var("b", "int64")
|
|
x0 = relax.Var("x", R.Tensor((a, b), "float32"))
|
|
|
|
tiling_transform = lambda a, b: (a, b // 3, b % 3)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.layout_transform(x0, index_map=tiling_transform),
|
|
relax.TensorType((a, (b - b % (-3)) // 3, 3), "float32"),
|
|
)
|
|
|
|
|
|
def test_layout_transform_infer_ty_shape_var():
|
|
bb = relax.BlockBuilder()
|
|
|
|
s = relax.Var("s", relax.ShapeType((30, 20)))
|
|
x = relax.Var("x", relax.TensorType(s, "float32"))
|
|
tiling_padding_transform = lambda a, b: (a, b // 3, b % 3)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.layout_transform(x, index_map=tiling_padding_transform),
|
|
relax.TensorType((30, 7, 3), "float32"),
|
|
)
|
|
|
|
s_unknown_shape = relax.Var("s", relax.ShapeType(ndim=2))
|
|
x_unknown_shape = relax.Var("x", relax.TensorType(s_unknown_shape, "float32"))
|
|
_check_inference(
|
|
bb,
|
|
relax.op.layout_transform(x_unknown_shape, index_map=tiling_padding_transform),
|
|
relax.TensorType(dtype="float32", ndim=3),
|
|
)
|
|
|
|
s_unknown_rank = relax.Var("s", relax.ShapeType())
|
|
x_unknown_rank = relax.Var("x", relax.TensorType(s_unknown_rank, "float32"))
|
|
_check_inference(
|
|
bb,
|
|
relax.op.layout_transform(x_unknown_rank, index_map=tiling_padding_transform),
|
|
relax.TensorType(dtype="float32", ndim=3),
|
|
)
|
|
|
|
a = tirx.Var("a", "int64")
|
|
b = tirx.Var("b", "int64")
|
|
s_symbolic_shape = relax.Var("s", relax.ShapeType((a, b)))
|
|
x_symbolic_shape = relax.Var("x", relax.TensorType(s_symbolic_shape, "float32"))
|
|
_check_inference(
|
|
bb,
|
|
relax.op.layout_transform(x_symbolic_shape, index_map=tiling_padding_transform),
|
|
relax.TensorType((a, (b - b % (-3)) // 3, 3), "float32"),
|
|
)
|
|
|
|
|
|
def test_layout_transform_infer_ty_invalid_index_map():
|
|
bb = relax.BlockBuilder()
|
|
x = relax.Var("x", R.Tensor((10, 20, 30), "float32"))
|
|
|
|
with pytest.raises(ValueError):
|
|
bb.normalize(relax.op.layout_transform(x, index_map=lambda a, b: (b, a)))
|
|
|
|
|
|
def test_squeeze_infer_ty():
|
|
bb = relax.BlockBuilder()
|
|
vdev0 = VDevice("llvm")
|
|
x0 = relax.Var("x", R.Tensor((2, 1, 3, 1, 1, 4), "float32"))
|
|
x1 = relax.Var("x", R.Tensor("float32", ndim=6))
|
|
x2 = relax.Var("x", R.Tensor("float32"))
|
|
x3 = relax.Var("x", R.Tensor((2, 1, 3, 1, 1, 4)))
|
|
x4 = relax.Var("x", R.Tensor(ndim=6))
|
|
x5 = relax.Var("x", R.Tensor())
|
|
x6 = relax.Var("x", R.Tensor((2, 1, 3, 1, 1, 4), "float32", vdev0))
|
|
|
|
_check_inference(bb, relax.op.squeeze(x0, [1, 4]), relax.TensorType((2, 3, 1, 4), "float32"))
|
|
_check_inference(
|
|
bb, relax.op.squeeze(x6, [1, 4]), relax.TensorType((2, 3, 1, 4), "float32", vdev0)
|
|
)
|
|
_check_inference(bb, relax.op.squeeze(x0), relax.TensorType((2, 3, 4), "float32"))
|
|
_check_inference(bb, relax.op.squeeze(x1, [1, 4]), relax.TensorType(dtype="float32", ndim=4))
|
|
_check_inference(bb, relax.op.squeeze(x1), relax.TensorType(dtype="float32"))
|
|
_check_inference(bb, relax.op.squeeze(x2, [1, 4]), relax.TensorType(dtype="float32"))
|
|
_check_inference(bb, relax.op.squeeze(x2), relax.TensorType(dtype="float32"))
|
|
_check_inference(bb, relax.op.squeeze(x3, [1, 4]), relax.TensorType((2, 3, 1, 4), dtype=""))
|
|
_check_inference(bb, relax.op.squeeze(x3), relax.TensorType((2, 3, 4), dtype=""))
|
|
_check_inference(bb, relax.op.squeeze(x4, [1, 4]), relax.TensorType(dtype="", ndim=4))
|
|
_check_inference(bb, relax.op.squeeze(x4), relax.TensorType(dtype=""))
|
|
_check_inference(bb, relax.op.squeeze(x5, [1, 4]), relax.TensorType(dtype=""))
|
|
_check_inference(bb, relax.op.squeeze(x5), relax.TensorType(dtype=""))
|
|
|
|
|
|
def test_squeeze_infer_ty_shape_symbolic():
|
|
bb = relax.BlockBuilder()
|
|
a = tirx.Var("a", "int64")
|
|
b = tirx.Var("b", "int64")
|
|
x0 = relax.Var("x", R.Tensor((a, 1, b), "float32"))
|
|
x1 = relax.Var("x", R.Tensor((a, 1, b)))
|
|
|
|
_check_inference(bb, relax.op.squeeze(x0, [1]), relax.TensorType((a, b), "float32"))
|
|
_check_inference(bb, relax.op.squeeze(x0), relax.TensorType(dtype="float32"))
|
|
_check_inference(bb, relax.op.squeeze(x1, [1]), relax.TensorType((a, b), dtype=""))
|
|
_check_inference(bb, relax.op.squeeze(x1), relax.TensorType(dtype=""))
|
|
|
|
|
|
def test_squeeze_infer_ty_shape_var():
|
|
bb = relax.BlockBuilder()
|
|
a = tirx.Var("a", "int64")
|
|
b = tirx.Var("b", "int64")
|
|
s0 = relax.Var("s", relax.ShapeType((2, 1, 3, 1, 1, 4)))
|
|
s1 = relax.Var("s", relax.ShapeType((2, 3, 4)))
|
|
s2 = relax.Var("s", relax.ShapeType((a, 1, b)))
|
|
s3 = relax.Var("s", relax.ShapeType(ndim=6))
|
|
s4 = relax.Var("s", relax.ShapeType())
|
|
x0 = relax.Var("x", relax.TensorType(s0, "float32"))
|
|
x1 = relax.Var("x", relax.TensorType(s1, "float32"))
|
|
x2 = relax.Var("x", relax.TensorType(s2, "float32"))
|
|
x3 = relax.Var("x", relax.TensorType(s3, "float32"))
|
|
x4 = relax.Var("x", relax.TensorType(s4, "float32"))
|
|
|
|
_check_inference(bb, relax.op.squeeze(x0, [1, 4]), relax.TensorType(dtype="float32", ndim=4))
|
|
_check_inference(bb, relax.op.squeeze(x0, []), relax.TensorType(s0, "float32"))
|
|
_check_inference(bb, relax.op.squeeze(x0), relax.TensorType(dtype="float32"))
|
|
_check_inference(bb, relax.op.squeeze(x1, []), relax.TensorType(s1, "float32"))
|
|
_check_inference(bb, relax.op.squeeze(x1), relax.TensorType(s1, dtype="float32"))
|
|
_check_inference(bb, relax.op.squeeze(x2, [1]), relax.TensorType(dtype="float32", ndim=2))
|
|
_check_inference(bb, relax.op.squeeze(x2, []), relax.TensorType(s2, "float32"))
|
|
_check_inference(bb, relax.op.squeeze(x2), relax.TensorType(dtype="float32"))
|
|
_check_inference(bb, relax.op.squeeze(x3, [1, 4]), relax.TensorType(dtype="float32", ndim=4))
|
|
_check_inference(bb, relax.op.squeeze(x3, []), relax.TensorType(s3, "float32"))
|
|
_check_inference(bb, relax.op.squeeze(x3), relax.TensorType(dtype="float32"))
|
|
_check_inference(bb, relax.op.squeeze(x4, [1, 4]), relax.TensorType(dtype="float32"))
|
|
_check_inference(bb, relax.op.squeeze(x4, []), relax.TensorType(s4, "float32"))
|
|
_check_inference(bb, relax.op.squeeze(x4), relax.TensorType(dtype="float32"))
|
|
|
|
|
|
def test_squeeze_infer_ty_more_input_dtype():
|
|
bb = relax.BlockBuilder()
|
|
x0 = relax.Var("x", R.Tensor((2, 1, 3, 1, 1, 4), "float16"))
|
|
x1 = relax.Var("x", R.Tensor((2, 1, 3, 1, 1, 4), "int8"))
|
|
x2 = relax.Var("x", R.Tensor((2, 1, 3, 1, 1, 4), "int32"))
|
|
|
|
_check_inference(bb, relax.op.squeeze(x0), relax.TensorType((2, 3, 4), "float16"))
|
|
_check_inference(bb, relax.op.squeeze(x1), relax.TensorType((2, 3, 4), "int8"))
|
|
_check_inference(bb, relax.op.squeeze(x2), relax.TensorType((2, 3, 4), "int32"))
|
|
|
|
|
|
def test_squeeze_infer_ty_axis_out_of_range():
|
|
bb = relax.BlockBuilder()
|
|
s0 = relax.Var("s", relax.ShapeType((2, 1, 3, 1, 1, 4)))
|
|
s1 = relax.Var("s", relax.ShapeType(ndim=6))
|
|
x0 = relax.Var("x", R.Tensor((2, 1, 3, 1, 1, 4), "float32"))
|
|
x1 = relax.Var("x", R.Tensor("float32", ndim=6))
|
|
x2 = relax.Var("x", relax.TensorType(s0, "float32"))
|
|
x3 = relax.Var("x", relax.TensorType(s1, "float32"))
|
|
|
|
with pytest.raises(ValueError):
|
|
bb.normalize(relax.op.squeeze(x0, [6]))
|
|
with pytest.raises(ValueError):
|
|
bb.normalize(relax.op.squeeze(x0, [-7]))
|
|
with pytest.raises(ValueError):
|
|
bb.normalize(relax.op.squeeze(x1, [6]))
|
|
with pytest.raises(ValueError):
|
|
bb.normalize(relax.op.squeeze(x1, [-7]))
|
|
with pytest.raises(ValueError):
|
|
bb.normalize(relax.op.squeeze(x2, [6]))
|
|
with pytest.raises(ValueError):
|
|
bb.normalize(relax.op.squeeze(x2, [-7]))
|
|
with pytest.raises(ValueError):
|
|
bb.normalize(relax.op.squeeze(x3, [6]))
|
|
with pytest.raises(ValueError):
|
|
bb.normalize(relax.op.squeeze(x3, [-7]))
|
|
|
|
|
|
def test_squeeze_infer_ty_repetitive_axes():
|
|
bb = relax.BlockBuilder()
|
|
s0 = relax.Var("s", relax.ShapeType((2, 1, 3, 1, 1, 4)))
|
|
s1 = relax.Var("s", relax.ShapeType(ndim=6))
|
|
x0 = relax.Var("x", R.Tensor((2, 1, 3, 1, 1, 4), "float32"))
|
|
x1 = relax.Var("x", R.Tensor("float32", ndim=6))
|
|
x2 = relax.Var("x", relax.TensorType(s0, "float32"))
|
|
x3 = relax.Var("x", relax.TensorType(s1, "float32"))
|
|
|
|
with pytest.raises(ValueError):
|
|
bb.normalize(relax.op.squeeze(x0, [3, -3]))
|
|
with pytest.raises(ValueError):
|
|
bb.normalize(relax.op.squeeze(x0, [1, 1]))
|
|
with pytest.raises(ValueError):
|
|
bb.normalize(relax.op.squeeze(x1, [3, -3]))
|
|
with pytest.raises(ValueError):
|
|
bb.normalize(relax.op.squeeze(x1, [1, 1]))
|
|
with pytest.raises(ValueError):
|
|
bb.normalize(relax.op.squeeze(x2, [3, -3]))
|
|
with pytest.raises(ValueError):
|
|
bb.normalize(relax.op.squeeze(x2, [1, 1]))
|
|
with pytest.raises(ValueError):
|
|
bb.normalize(relax.op.squeeze(x3, [3, -3]))
|
|
with pytest.raises(ValueError):
|
|
bb.normalize(relax.op.squeeze(x3, [1, 1]))
|
|
|
|
|
|
def test_squeeze_infer_ty_axis_length_not_one():
|
|
bb = relax.BlockBuilder()
|
|
a = tirx.Var("a", "int64")
|
|
s0 = relax.Var("s", relax.ShapeType((2, 3, 4)))
|
|
s1 = relax.Var("s", relax.ShapeType((a, 3, 4)))
|
|
x0 = relax.Var("x", R.Tensor((2, 3, 4), "float32"))
|
|
x1 = relax.Var("x", R.Tensor((a, 3, 4), "float32"))
|
|
x2 = relax.Var("x", relax.TensorType(s0, "float32"))
|
|
x3 = relax.Var("x", relax.TensorType(s1, "float32"))
|
|
|
|
# Squeeze concrete shape (2,3,4) at axis=0, but axis length 2 != 1, squeeze is no-op.
|
|
_check_inference(
|
|
bb, relax.op.squeeze(x0, [0]), relax.TensorType(shape=(2, 3, 4), dtype="float32")
|
|
)
|
|
# Squeeze symbolic shape (a,3,4) at axis=0, assuming a can achieve successful squeeze.
|
|
_check_inference(bb, relax.op.squeeze(x1, [0]), relax.TensorType(shape=(3, 4), dtype="float32"))
|
|
# Squeeze shape variable s0 (corresponding to (2,3,4)) at axis=0.
|
|
_check_inference(bb, relax.op.squeeze(x2, [0]), relax.TensorType(shape=s0, dtype="float32"))
|
|
# Squeeze shape variable s1 (a,3,4) at axis=0, assuming a can achieve successful squeeze.
|
|
_check_inference(bb, relax.op.squeeze(x3, [0]), relax.TensorType(dtype="float32", ndim=2))
|
|
|
|
|
|
def test_squeeze_infer_ty_wrong_input_type():
|
|
bb = relax.BlockBuilder()
|
|
x0 = relax.Var("x", relax.ShapeType((2, 3, 4)))
|
|
x1 = relax.Var("x", relax.FuncType([], R.Tensor((2, 3, 4), "float32")))
|
|
|
|
with pytest.raises(TypeError):
|
|
bb.normalize(relax.op.squeeze(x0))
|
|
with pytest.raises(TypeError):
|
|
bb.normalize(relax.op.squeeze(x1))
|
|
|
|
|
|
def test_flatten_infer_ty():
|
|
bb = relax.BlockBuilder()
|
|
vdev0 = VDevice("llvm")
|
|
x0 = relax.Var("x", R.Tensor((3, 4, 5), "float32"))
|
|
x1 = relax.Var("x", R.Tensor((3,), "float32"))
|
|
x2 = relax.Var("x", R.Tensor((), "float32"))
|
|
x3 = relax.Var("x", R.Tensor("float32", ndim=3))
|
|
x4 = relax.Var("x", R.Tensor("float32", ndim=1))
|
|
x5 = relax.Var("x", R.Tensor("float32", ndim=0))
|
|
x6 = relax.Var("x", R.Tensor("float32"))
|
|
x7 = relax.Var("x", R.Tensor((3, 4, 5)))
|
|
x8 = relax.Var("x", R.Tensor((3,)))
|
|
x9 = relax.Var("x", R.Tensor(()))
|
|
x10 = relax.Var("x", R.Tensor(ndim=3))
|
|
x11 = relax.Var("x", R.Tensor(ndim=1))
|
|
x12 = relax.Var("x", R.Tensor(ndim=0))
|
|
x13 = relax.Var("x", R.Tensor())
|
|
x14 = relax.Var("x", R.Tensor((3, 4, 5), "float32", vdev0))
|
|
|
|
_check_inference(bb, relax.op.flatten(x0), relax.TensorType((60,), "float32"))
|
|
_check_inference(bb, relax.op.flatten(x14), relax.TensorType((60,), "float32", vdev0))
|
|
_check_inference(bb, relax.op.flatten(x1), relax.TensorType((3,), "float32"))
|
|
_check_inference(bb, relax.op.flatten(x2), relax.TensorType((1,), "float32"))
|
|
_check_inference(bb, relax.op.flatten(x3), relax.TensorType(dtype="float32", ndim=1))
|
|
_check_inference(bb, relax.op.flatten(x4), relax.TensorType(dtype="float32", ndim=1))
|
|
_check_inference(bb, relax.op.flatten(x5), relax.TensorType((1,), "float32"))
|
|
_check_inference(bb, relax.op.flatten(x6), relax.TensorType(dtype="float32", ndim=1))
|
|
_check_inference(bb, relax.op.flatten(x7), relax.TensorType((60,), dtype=""))
|
|
_check_inference(bb, relax.op.flatten(x8), relax.TensorType((3,), dtype=""))
|
|
_check_inference(bb, relax.op.flatten(x9), relax.TensorType((1,), dtype=""))
|
|
_check_inference(bb, relax.op.flatten(x10), relax.TensorType(dtype="", ndim=1))
|
|
_check_inference(bb, relax.op.flatten(x11), relax.TensorType(dtype="", ndim=1))
|
|
_check_inference(bb, relax.op.flatten(x12), relax.TensorType((1,), dtype=""))
|
|
_check_inference(bb, relax.op.flatten(x13), relax.TensorType(dtype="", ndim=1))
|
|
|
|
|
|
def test_flatten_infer_ty_shape_symbolic():
|
|
bb = relax.BlockBuilder()
|
|
a = tirx.Var("a", "int64")
|
|
b = tirx.Var("b", "int64")
|
|
x0 = relax.Var("x", R.Tensor((a, b), "float32"))
|
|
x1 = relax.Var("x", R.Tensor((a, b)))
|
|
|
|
_check_inference(bb, relax.op.flatten(x0), relax.TensorType((a * b,), "float32"))
|
|
_check_inference(bb, relax.op.flatten(x1), relax.TensorType((a * b,), dtype=""))
|
|
|
|
|
|
def test_flatten_infer_ty_shape_var():
|
|
bb = relax.BlockBuilder()
|
|
s0 = relax.Var("s", relax.ShapeType((3, 4, 5)))
|
|
s1 = relax.Var("s", relax.ShapeType((3,)))
|
|
s2 = relax.Var("s", relax.ShapeType(()))
|
|
s3 = relax.Var("s", relax.ShapeType(ndim=3))
|
|
s4 = relax.Var("s", relax.ShapeType(ndim=1))
|
|
s5 = relax.Var("s", relax.ShapeType(ndim=0))
|
|
s6 = relax.Var("s", relax.ShapeType())
|
|
x0 = relax.Var("x", relax.TensorType(s0, "float32"))
|
|
x1 = relax.Var("x", relax.TensorType(s1, "float32"))
|
|
x2 = relax.Var("x", relax.TensorType(s2, "float32"))
|
|
x3 = relax.Var("x", relax.TensorType(s3, "float32"))
|
|
x4 = relax.Var("x", relax.TensorType(s4, "float32"))
|
|
x5 = relax.Var("x", relax.TensorType(s5, "float32"))
|
|
x6 = relax.Var("x", relax.TensorType(s6, "float32"))
|
|
|
|
_check_inference(bb, relax.op.flatten(x0), relax.TensorType(dtype="float32", ndim=1))
|
|
_check_inference(bb, relax.op.flatten(x1), relax.TensorType(s1, "float32"))
|
|
_check_inference(bb, relax.op.flatten(x2), relax.TensorType((1,), "float32"))
|
|
_check_inference(bb, relax.op.flatten(x3), relax.TensorType(dtype="float32", ndim=1))
|
|
_check_inference(bb, relax.op.flatten(x4), relax.TensorType(s4, "float32"))
|
|
_check_inference(bb, relax.op.flatten(x5), relax.TensorType((1,), "float32"))
|
|
_check_inference(bb, relax.op.flatten(x6), relax.TensorType(dtype="float32", ndim=1))
|
|
|
|
|
|
def test_flatten_infer_ty_more_input_dtype():
|
|
bb = relax.BlockBuilder()
|
|
x0 = relax.Var("x", R.Tensor((3, 4, 5), "float16"))
|
|
x1 = relax.Var("x", R.Tensor((3, 4, 5), "int8"))
|
|
x2 = relax.Var("x", R.Tensor((3, 4, 5), "int32"))
|
|
|
|
_check_inference(bb, relax.op.flatten(x0), relax.TensorType((60,), "float16"))
|
|
_check_inference(bb, relax.op.flatten(x1), relax.TensorType((60,), "int8"))
|
|
_check_inference(bb, relax.op.flatten(x2), relax.TensorType((60,), "int32"))
|
|
|
|
|
|
def test_flatten_infer_ty_wrong_input_type():
|
|
bb = relax.BlockBuilder()
|
|
x0 = relax.Var("x", relax.ShapeType((3, 4, 5)))
|
|
x1 = relax.Var("x", relax.FuncType([], R.Tensor((3, 4, 5), "float32")))
|
|
|
|
with pytest.raises(TypeError):
|
|
bb.normalize(relax.op.flatten(x0))
|
|
with pytest.raises(TypeError):
|
|
bb.normalize(relax.op.flatten(x1))
|
|
|
|
|
|
def test_flatten_wrong_input_number():
|
|
x = relax.Var("x", R.Tensor((3, 4, 5), "float32"))
|
|
y = relax.Var("y", R.Tensor((2, 3, 4), "float32"))
|
|
|
|
with pytest.raises(TypeError):
|
|
relax.op.flatten(x, y)
|
|
|
|
|
|
def test_concat_infer_ty_with_axis():
|
|
bb = relax.BlockBuilder()
|
|
vdev0 = VDevice("llvm")
|
|
x0 = relax.Var("x", R.Tensor((2, 3, 4), "float32"))
|
|
x1 = relax.Var("x", R.Tensor("float32", ndim=3))
|
|
x2 = relax.Var("x", R.Tensor("float32"))
|
|
x3 = relax.Var("x", R.Tensor((2, 3, 4)))
|
|
x4 = relax.Var("x", R.Tensor(ndim=3))
|
|
x5 = relax.Var("x", R.Tensor())
|
|
x6 = relax.Var("x", R.Tensor((2, 3, 4), "float32", vdev0))
|
|
y0 = relax.Var("y", R.Tensor((2, 4, 4), "float32"))
|
|
y1 = relax.Var("y", R.Tensor("float32", ndim=3))
|
|
y2 = relax.Var("y", R.Tensor("float32"))
|
|
y3 = relax.Var("y", R.Tensor((2, 4, 4)))
|
|
y4 = relax.Var("y", R.Tensor(ndim=3))
|
|
y5 = relax.Var("y", R.Tensor())
|
|
y6 = relax.Var("y", R.Tensor((2, 4, 4), "float32", vdev0))
|
|
z0 = relax.Var("z", R.Tensor((2, 5, 4), "float32"))
|
|
z1 = relax.Var("z", R.Tensor("float32", ndim=3))
|
|
z2 = relax.Var("z", R.Tensor("float32"))
|
|
z3 = relax.Var("z", R.Tensor((2, 5, 4)))
|
|
z4 = relax.Var("z", R.Tensor(ndim=3))
|
|
z5 = relax.Var("z", R.Tensor())
|
|
z6 = relax.Var("z", R.Tensor((2, 5, 4), "float32", vdev0))
|
|
|
|
_check_inference(
|
|
bb, relax.op.concat([x0, y0, z0], axis=1), relax.TensorType((2, 12, 4), "float32")
|
|
)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.concat([x6, y6, z6], axis=1),
|
|
relax.TensorType((2, 12, 4), "float32", vdev0),
|
|
)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.concat([x6, y0, z0], axis=1),
|
|
relax.TensorType((2, 12, 4), "float32", vdev0),
|
|
)
|
|
_check_inference(
|
|
bb, relax.op.concat([x0, y0, z0], axis=-2), relax.TensorType((2, 12, 4), "float32")
|
|
)
|
|
_check_inference(
|
|
bb, relax.op.concat([x1, y0, z0], axis=1), relax.TensorType(dtype="float32", ndim=3)
|
|
)
|
|
_check_inference(
|
|
bb, relax.op.concat([x2, y0, z0], axis=1), relax.TensorType(dtype="float32", ndim=3)
|
|
)
|
|
_check_inference(
|
|
bb, relax.op.concat([x3, y0, z0], axis=1), relax.TensorType((2, 12, 4), dtype="")
|
|
)
|
|
_check_inference(
|
|
bb, relax.op.concat([x3, y0, z0], axis=-2), relax.TensorType((2, 12, 4), dtype="")
|
|
)
|
|
_check_inference(bb, relax.op.concat([x4, y0, z0], axis=1), relax.TensorType(dtype="", ndim=3))
|
|
_check_inference(bb, relax.op.concat([x5, y0, z0], axis=1), relax.TensorType(dtype="", ndim=3))
|
|
_check_inference(
|
|
bb, relax.op.concat([x1, y1, z0], axis=1), relax.TensorType(dtype="float32", ndim=3)
|
|
)
|
|
_check_inference(
|
|
bb, relax.op.concat([x2, y1, z0], axis=1), relax.TensorType(dtype="float32", ndim=3)
|
|
)
|
|
_check_inference(bb, relax.op.concat([x3, y1, z0], axis=1), relax.TensorType(dtype="", ndim=3))
|
|
_check_inference(bb, relax.op.concat([x5, y1, z0], axis=1), relax.TensorType(dtype="", ndim=3))
|
|
_check_inference(
|
|
bb, relax.op.concat([x2, y2, z0], axis=1), relax.TensorType(dtype="float32", ndim=3)
|
|
)
|
|
_check_inference(bb, relax.op.concat([x3, y2, z0], axis=1), relax.TensorType(dtype="", ndim=3))
|
|
_check_inference(bb, relax.op.concat([x5, y5, z0], axis=1), relax.TensorType(dtype="", ndim=3))
|
|
_check_inference(
|
|
bb, relax.op.concat([x1, y1, z1], axis=1), relax.TensorType(dtype="float32", ndim=3)
|
|
)
|
|
_check_inference(
|
|
bb, relax.op.concat([x2, y2, z1], axis=1), relax.TensorType(dtype="float32", ndim=3)
|
|
)
|
|
_check_inference(bb, relax.op.concat([x3, y1, z1], axis=1), relax.TensorType(dtype="", ndim=3))
|
|
_check_inference(
|
|
bb, relax.op.concat([x2, y2, z2], axis=1), relax.TensorType(dtype="float32", ndim=-1)
|
|
)
|
|
_check_inference(bb, relax.op.concat([x3, y2, z2], axis=1), relax.TensorType(dtype="", ndim=3))
|
|
_check_inference(bb, relax.op.concat([x4, y4, z2], axis=1), relax.TensorType(dtype="", ndim=3))
|
|
_check_inference(bb, relax.op.concat([x5, y5, z2], axis=1), relax.TensorType(dtype="", ndim=-1))
|
|
_check_inference(
|
|
bb, relax.op.concat([x3, y3, z3], axis=1), relax.TensorType((2, 12, 4), dtype="")
|
|
)
|
|
_check_inference(
|
|
bb, relax.op.concat([x3, y3, z3], axis=-2), relax.TensorType((2, 12, 4), dtype="")
|
|
)
|
|
_check_inference(bb, relax.op.concat([x4, y3, z3], axis=1), relax.TensorType(dtype="", ndim=3))
|
|
_check_inference(bb, relax.op.concat([x5, y5, z3], axis=1), relax.TensorType(dtype="", ndim=3))
|
|
_check_inference(bb, relax.op.concat([x4, y4, z4], axis=1), relax.TensorType(dtype="", ndim=3))
|
|
_check_inference(bb, relax.op.concat([x5, y5, z4], axis=1), relax.TensorType(dtype="", ndim=3))
|
|
_check_inference(bb, relax.op.concat([x5, y5, z5], axis=1), relax.TensorType(dtype=""))
|
|
_check_inference(
|
|
bb,
|
|
relax.op.concat(relax.Tuple([x0, y0, z0]), axis=1),
|
|
relax.TensorType((2, 12, 4), "float32"),
|
|
)
|
|
|
|
|
|
def test_concat_infer_ty_with_axis_shape_symbolic():
|
|
bb = relax.BlockBuilder()
|
|
a0 = tirx.Var("a0", "int64")
|
|
a1 = tirx.Var("a1", "int64")
|
|
b0 = tirx.Var("b0", "int64")
|
|
b1 = tirx.Var("b1", "int64")
|
|
b2 = tirx.Var("b2", "int64")
|
|
c = tirx.Var("c", "int64")
|
|
x0 = relax.Var("x", R.Tensor((a0, b0, c), "float32"))
|
|
x1 = relax.Var("x", R.Tensor((a1, b0, c), "float32"))
|
|
x2 = relax.Var("x", R.Tensor((a0, b0, c), "float32"))
|
|
y = relax.Var("y", R.Tensor((a0, b1, c), "float32"))
|
|
z = relax.Var("z", R.Tensor((a0, b2, c), "float32"))
|
|
|
|
_check_inference(
|
|
bb,
|
|
relax.op.concat([x0, y, z], axis=1),
|
|
relax.TensorType((a0, b0 + b1 + b2, c), "float32"),
|
|
)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.concat([x0, y, z], axis=-2),
|
|
relax.TensorType((a0, b0 + b1 + b2, c), "float32"),
|
|
)
|
|
_check_inference(
|
|
bb, relax.op.concat([x1, y, z], axis=1), relax.TensorType(dtype="float32", ndim=3)
|
|
)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.concat(relax.Tuple([x0, y, z]), axis=1),
|
|
relax.TensorType((a0, b0 + b1 + b2, c), "float32"),
|
|
)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.concat(relax.Tuple([x0, x2]), axis=1),
|
|
relax.TensorType((a0, b0 * 2, c), "float32"),
|
|
)
|
|
|
|
|
|
def test_concat_infer_ty_with_axis_shape_var():
|
|
bb = relax.BlockBuilder()
|
|
a0 = tirx.Var("a0", "int64")
|
|
a1 = tirx.Var("a1", "int64")
|
|
b0 = tirx.Var("b0", "int64")
|
|
b1 = tirx.Var("b1", "int64")
|
|
b2 = tirx.Var("b2", "int64")
|
|
c = tirx.Var("c", "int64")
|
|
sx0 = relax.Var("sx", relax.ShapeType((2, 3, 4)))
|
|
sx1 = relax.Var("sx", relax.ShapeType((a0, b0, c)))
|
|
sx2 = relax.Var("sx", relax.ShapeType((a1, b0, c)))
|
|
sx3 = relax.Var("sx", relax.ShapeType(ndim=3))
|
|
sx4 = relax.Var("sx", relax.ShapeType())
|
|
x0 = relax.Var("x", relax.TensorType(sx0, "float32"))
|
|
x1 = relax.Var("x", relax.TensorType(sx1, "float32"))
|
|
x2 = relax.Var("x", relax.TensorType(sx2, "float32"))
|
|
x3 = relax.Var("x", relax.TensorType(sx3, "float32"))
|
|
x4 = relax.Var("x", relax.TensorType(sx4, "float32"))
|
|
y0 = relax.Var("y", R.Tensor((2, 4, 4), "float32"))
|
|
y1 = relax.Var("y", R.Tensor((a0, b1, c), "float32"))
|
|
z0 = relax.Var("z", R.Tensor((2, 5, 4), "float32"))
|
|
z1 = relax.Var("z", R.Tensor((a0, b2, c), "float32"))
|
|
|
|
_check_inference(
|
|
bb, relax.op.concat([x0, y0, z0], axis=1), relax.TensorType(dtype="float32", ndim=3)
|
|
)
|
|
_check_inference(
|
|
bb, relax.op.concat([x1, y1, z1], axis=1), relax.TensorType(dtype="float32", ndim=3)
|
|
)
|
|
_check_inference(
|
|
bb, relax.op.concat([x2, y1, z1], axis=1), relax.TensorType(dtype="float32", ndim=3)
|
|
)
|
|
_check_inference(
|
|
bb, relax.op.concat([x3, y0, z0], axis=1), relax.TensorType(dtype="float32", ndim=3)
|
|
)
|
|
_check_inference(
|
|
bb, relax.op.concat([x4, y0, z0], axis=1), relax.TensorType(dtype="float32", ndim=3)
|
|
)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.concat(relax.Tuple([x0, y0, z0]), axis=1),
|
|
relax.TensorType(dtype="float32", ndim=3),
|
|
)
|
|
|
|
|
|
def test_concat_infer_ty_without_axis():
|
|
bb = relax.BlockBuilder()
|
|
x0 = relax.Var("x", R.Tensor((3,), "float32"))
|
|
x1 = relax.Var("x", R.Tensor("float32", ndim=1))
|
|
x2 = relax.Var("x", R.Tensor((3,)))
|
|
x3 = relax.Var("x", R.Tensor(ndim=1))
|
|
y0 = relax.Var("y", R.Tensor((4,), "float32"))
|
|
y1 = relax.Var("y", R.Tensor("float32", ndim=1))
|
|
z0 = relax.Var("z", R.Tensor((5,), "float32"))
|
|
z1 = relax.Var("z", R.Tensor("float32", ndim=1))
|
|
|
|
_check_inference(
|
|
bb, relax.op.concat([x0, y0, z0], axis=None), relax.TensorType((12,), "float32")
|
|
)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.concat([x1, y0, z0], axis=None),
|
|
relax.TensorType(dtype="float32", ndim=1),
|
|
)
|
|
_check_inference(
|
|
bb, relax.op.concat([x2, y0, z0], axis=None), relax.TensorType((12,), dtype="")
|
|
)
|
|
_check_inference(
|
|
bb, relax.op.concat([x3, y0, z0], axis=None), relax.TensorType(dtype="", ndim=1)
|
|
)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.concat([x1, y1, z0], axis=None),
|
|
relax.TensorType(dtype="float32", ndim=1),
|
|
)
|
|
_check_inference(
|
|
bb, relax.op.concat([x2, y1, z0], axis=None), relax.TensorType(dtype="", ndim=1)
|
|
)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.concat([x1, y1, z1], axis=None),
|
|
relax.TensorType(dtype="float32", ndim=1),
|
|
)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.concat(relax.Tuple([x0, y0, z0]), axis=None),
|
|
relax.TensorType((12,), "float32"),
|
|
)
|
|
|
|
|
|
def test_concat_infer_ty_without_axis_shape_symbolic():
|
|
bb = relax.BlockBuilder()
|
|
a0 = tirx.Var("a0", "int64")
|
|
a1 = tirx.Var("a1", "int64")
|
|
x0 = relax.Var("x", R.Tensor((a0,), "float32"))
|
|
x1 = relax.Var("x", R.Tensor((a0,), ""))
|
|
y0 = relax.Var("y", R.Tensor((a1,), "float32"))
|
|
y1 = relax.Var("y", R.Tensor((a1,), ""))
|
|
|
|
_check_inference(
|
|
bb, relax.op.concat([x0, y0], axis=None), relax.TensorType((a0 + a1,), "float32")
|
|
)
|
|
_check_inference(
|
|
bb, relax.op.concat([x0, y1], axis=None), relax.TensorType((a0 + a1,), dtype="")
|
|
)
|
|
_check_inference(
|
|
bb, relax.op.concat([x1, y0], axis=None), relax.TensorType((a0 + a1,), dtype="")
|
|
)
|
|
_check_inference(
|
|
bb, relax.op.concat([x1, y1], axis=None), relax.TensorType((a0 + a1,), dtype="")
|
|
)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.concat(relax.Tuple([x0, y0]), axis=None),
|
|
relax.TensorType((a0 + a1,), "float32"),
|
|
)
|
|
|
|
|
|
def test_concat_infer_ty_without_axis_shape_var():
|
|
bb = relax.BlockBuilder()
|
|
sx0 = relax.Var("sx", relax.ShapeType((3,)))
|
|
sx1 = relax.Var("sx", relax.ShapeType(ndim=1))
|
|
sy0 = relax.Var("sy", relax.ShapeType((4,)))
|
|
x0 = relax.Var("x", relax.TensorType(sx0, "float32"))
|
|
x1 = relax.Var("x", relax.TensorType(sx1, "float32"))
|
|
y0 = relax.Var("y", relax.TensorType(sy0, "float32"))
|
|
|
|
_check_inference(
|
|
bb, relax.op.concat([x0, y0], axis=None), relax.TensorType(dtype="float32", ndim=1)
|
|
)
|
|
_check_inference(
|
|
bb, relax.op.concat([x1, y0], axis=None), relax.TensorType(dtype="float32", ndim=1)
|
|
)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.concat(relax.Tuple([x0, y0]), axis=None),
|
|
relax.TensorType(dtype="float32", ndim=1),
|
|
)
|
|
|
|
|
|
def test_concat_infer_ty_more_input_dtype():
|
|
bb = relax.BlockBuilder()
|
|
x0 = relax.Var("x", R.Tensor((3,), "float16"))
|
|
y0 = relax.Var("y", R.Tensor((4,), "float16"))
|
|
x1 = relax.Var("x", R.Tensor((3,), "int8"))
|
|
y1 = relax.Var("y", R.Tensor((4,), "int8"))
|
|
x2 = relax.Var("x", R.Tensor((3,), "int32"))
|
|
y2 = relax.Var("y", R.Tensor((4,), "int32"))
|
|
|
|
_check_inference(bb, relax.op.concat([x0, y0], axis=None), relax.TensorType((7,), "float16"))
|
|
_check_inference(bb, relax.op.concat([x1, y1], axis=None), relax.TensorType((7,), "int8"))
|
|
_check_inference(bb, relax.op.concat([x2, y2], axis=None), relax.TensorType((7,), "int32"))
|
|
|
|
|
|
def test_concat_infer_ty_tuple_var():
|
|
bb = relax.BlockBuilder()
|
|
a = tirx.Var("a0", "int64")
|
|
b0 = tirx.Var("b0", "int64")
|
|
b1 = tirx.Var("b1", "int64")
|
|
t0 = relax.Var(
|
|
"t",
|
|
relax.TupleType(
|
|
[relax.TensorType((a, b0), "float32"), relax.TensorType((a, b1), "float32")]
|
|
),
|
|
)
|
|
t1 = relax.Var(
|
|
"t",
|
|
relax.TupleType(
|
|
[
|
|
relax.TensorType((a, b0), "float32"),
|
|
relax.TensorType(dtype="float32", ndim=2),
|
|
]
|
|
),
|
|
)
|
|
t2 = relax.Var(
|
|
"t",
|
|
relax.TupleType(
|
|
[
|
|
relax.TensorType(dtype="float32"),
|
|
relax.TensorType(dtype="float32", ndim=2),
|
|
]
|
|
),
|
|
)
|
|
t3 = relax.Var(
|
|
"t",
|
|
relax.TupleType([relax.TensorType(dtype="float32"), relax.TensorType(dtype="float32")]),
|
|
)
|
|
t4 = relax.Var(
|
|
"t",
|
|
relax.TupleType([relax.TensorType((a, b0), "float32"), relax.TensorType((a, b1))]),
|
|
)
|
|
t5 = relax.Var(
|
|
"t",
|
|
relax.TupleType([relax.TensorType((a, b0), dtype=""), relax.TensorType((a, b1), dtype="")]),
|
|
)
|
|
t6 = relax.Var(
|
|
"t",
|
|
relax.TupleType([relax.TensorType(dtype="", ndim=2), relax.TensorType(dtype="")]),
|
|
)
|
|
t7 = relax.Var(
|
|
"t",
|
|
relax.TupleType([relax.TensorType(dtype=""), relax.TensorType(dtype="")]),
|
|
)
|
|
|
|
_check_inference(bb, relax.op.concat(t0, axis=1), relax.TensorType((a, b0 + b1), "float32"))
|
|
_check_inference(bb, relax.op.concat(t1, axis=1), relax.TensorType(dtype="float32", ndim=2))
|
|
_check_inference(bb, relax.op.concat(t2, axis=1), relax.TensorType(dtype="float32", ndim=2))
|
|
_check_inference(bb, relax.op.concat(t3, axis=1), relax.TensorType(dtype="float32"))
|
|
_check_inference(bb, relax.op.concat(t4, axis=1), relax.TensorType((a, b0 + b1), "float32"))
|
|
_check_inference(bb, relax.op.concat(t5, axis=1), relax.TensorType((a, b0 + b1), dtype=""))
|
|
_check_inference(bb, relax.op.concat(t6, axis=1), relax.TensorType(dtype="", ndim=2))
|
|
_check_inference(bb, relax.op.concat(t7, axis=1), relax.TensorType(dtype=""))
|
|
|
|
|
|
def test_concat_infer_ty_single_input_tensor():
|
|
bb = relax.BlockBuilder()
|
|
a = tirx.Var("a", "int64")
|
|
s0 = relax.Var("s", relax.ShapeType((3, a)))
|
|
s1 = relax.Var("s", relax.ShapeType((a,)))
|
|
s2 = relax.Var("s", relax.ShapeType(ndim=3))
|
|
s3 = relax.Var("s", relax.ShapeType(ndim=1))
|
|
s4 = relax.Var("s", relax.ShapeType())
|
|
x0 = relax.Var("x", R.Tensor((3, a), "float32"))
|
|
x1 = relax.Var("x", R.Tensor((a,), "float32"))
|
|
x2 = relax.Var("x", R.Tensor("float32", ndim=3))
|
|
x3 = relax.Var("x", R.Tensor("float32", ndim=1))
|
|
x4 = relax.Var("x", R.Tensor("float32"))
|
|
x5 = relax.Var("x", relax.TensorType(s0, "float32"))
|
|
x6 = relax.Var("x", relax.TensorType(s1, "float32"))
|
|
x7 = relax.Var("x", relax.TensorType(s2, "float32"))
|
|
x8 = relax.Var("x", relax.TensorType(s3, "float32"))
|
|
x9 = relax.Var("x", relax.TensorType(s4, "float32"))
|
|
|
|
_check_inference(bb, relax.op.concat([x0], axis=1), relax.TensorType((3, a), "float32"))
|
|
_check_inference(bb, relax.op.concat([x1], axis=0), relax.TensorType((a,), "float32"))
|
|
_check_inference(bb, relax.op.concat([x1], axis=None), relax.TensorType((a,), "float32"))
|
|
_check_inference(bb, relax.op.concat([x2], axis=1), relax.TensorType(dtype="float32", ndim=3))
|
|
_check_inference(bb, relax.op.concat([x3], axis=0), relax.TensorType(dtype="float32", ndim=1))
|
|
_check_inference(
|
|
bb, relax.op.concat([x3], axis=None), relax.TensorType(dtype="float32", ndim=1)
|
|
)
|
|
_check_inference(bb, relax.op.concat([x4], axis=1), relax.TensorType(dtype="float32"))
|
|
_check_inference(bb, relax.op.concat([x5], axis=1), relax.TensorType(s0, dtype="float32"))
|
|
_check_inference(bb, relax.op.concat([x6], axis=0), relax.TensorType(s1, dtype="float32"))
|
|
_check_inference(bb, relax.op.concat([x6], axis=None), relax.TensorType(s1, dtype="float32"))
|
|
_check_inference(bb, relax.op.concat([x7], axis=1), relax.TensorType(s2, dtype="float32"))
|
|
_check_inference(bb, relax.op.concat([x8], axis=0), relax.TensorType(s3, dtype="float32"))
|
|
_check_inference(bb, relax.op.concat([x8], axis=None), relax.TensorType(s3, dtype="float32"))
|
|
_check_inference(bb, relax.op.concat([x9], axis=1), relax.TensorType(s4, dtype="float32"))
|
|
|
|
|
|
def test_concat_infer_ty_zero_rank_input_tensor():
|
|
bb = relax.BlockBuilder()
|
|
s0 = relax.Var("s", relax.ShapeType(()))
|
|
s1 = relax.Var("s", relax.ShapeType(ndim=0))
|
|
x0 = relax.Var("x", R.Tensor((), "float32"))
|
|
x1 = relax.Var("x", R.Tensor("float32", ndim=0))
|
|
x2 = relax.Var("x", relax.TensorType(s0, "float32"))
|
|
x3 = relax.Var("x", relax.TensorType(s1, "float32"))
|
|
|
|
with pytest.raises(ValueError):
|
|
bb.normalize(relax.op.concat([x0], axis=0))
|
|
with pytest.raises(ValueError):
|
|
bb.normalize(relax.op.concat([x1], axis=0))
|
|
with pytest.raises(ValueError):
|
|
bb.normalize(relax.op.concat([x2], axis=None))
|
|
with pytest.raises(ValueError):
|
|
bb.normalize(relax.op.concat([x3], axis=None))
|
|
|
|
|
|
def test_concat_infer_ty_no_input_tensor():
|
|
bb = relax.BlockBuilder()
|
|
with pytest.raises(ValueError):
|
|
bb.normalize(relax.op.concat([], axis=1))
|
|
with pytest.raises(ValueError):
|
|
bb.normalize(relax.op.concat([], axis=None))
|
|
|
|
|
|
def test_concat_infer_ty_without_axis_but_tensor_not_one_dimensional():
|
|
bb = relax.BlockBuilder()
|
|
s0 = relax.Var("s", relax.ShapeType((3, 4)))
|
|
s1 = relax.Var("s", relax.ShapeType(ndim=2))
|
|
s2 = relax.Var("s", relax.ShapeType())
|
|
x0 = relax.Var("x", R.Tensor((3, 4), "float32"))
|
|
x1 = relax.Var("x", R.Tensor("float32", ndim=2))
|
|
x2 = relax.Var("x", R.Tensor("float32"))
|
|
x3 = relax.Var("x", relax.TensorType(s0, "float32"))
|
|
x4 = relax.Var("x", relax.TensorType(s1, "float32"))
|
|
x5 = relax.Var("x", relax.TensorType(s2, "float32"))
|
|
|
|
with pytest.raises(ValueError):
|
|
bb.normalize(relax.op.concat([x0], axis=None))
|
|
with pytest.raises(ValueError):
|
|
bb.normalize(relax.op.concat([x1], axis=None))
|
|
_check_inference(bb, relax.op.concat([x2], axis=None), relax.TensorType(dtype="float32"))
|
|
with pytest.raises(ValueError):
|
|
bb.normalize(relax.op.concat([x3], axis=None))
|
|
with pytest.raises(ValueError):
|
|
bb.normalize(relax.op.concat([x4], axis=None))
|
|
_check_inference(bb, relax.op.concat([x5], axis=None), relax.TensorType(s2, "float32"))
|
|
|
|
|
|
def test_concat_infer_ty_inconsistent_dtype():
|
|
bb = relax.BlockBuilder()
|
|
x = relax.Var("x", R.Tensor((3,)))
|
|
y = relax.Var("y", R.Tensor((4,), "float32"))
|
|
z = relax.Var("z", R.Tensor((5,), "int8"))
|
|
|
|
with pytest.raises(TypeError):
|
|
bb.normalize(relax.op.concat([x, y, z], axis=0))
|
|
|
|
|
|
def test_concat_infer_ty_inconsistent_ndim():
|
|
bb = relax.BlockBuilder()
|
|
s0 = relax.Var("s", relax.ShapeType((4, 5)))
|
|
s1 = relax.Var("s", relax.ShapeType(ndim=2))
|
|
x = relax.Var("x", R.Tensor((3,), "float32"))
|
|
y0 = relax.Var("y", R.Tensor((4, 5), "float32"))
|
|
y1 = relax.Var("y", R.Tensor("float32", ndim=2))
|
|
y2 = relax.Var("y", relax.TensorType(s0, "float32"))
|
|
y3 = relax.Var("y", relax.TensorType(s1, "float32"))
|
|
z = relax.Var("z", R.Tensor((5,), "float32"))
|
|
|
|
with pytest.raises(ValueError):
|
|
bb.normalize(relax.op.concat([x, y0, z], axis=0))
|
|
with pytest.raises(ValueError):
|
|
bb.normalize(relax.op.concat([x, y1, z], axis=0))
|
|
with pytest.raises(ValueError):
|
|
bb.normalize(relax.op.concat([x, y2, z], axis=0))
|
|
with pytest.raises(ValueError):
|
|
bb.normalize(relax.op.concat([x, y3, z], axis=0))
|
|
|
|
|
|
def test_concat_infer_ty_axis_out_of_range():
|
|
bb = relax.BlockBuilder()
|
|
s0 = relax.Var("s", relax.ShapeType((3,)))
|
|
s1 = relax.Var("s", relax.ShapeType(ndim=1))
|
|
x0 = relax.Var("x", R.Tensor((3,), "float32"))
|
|
x1 = relax.Var("x", R.Tensor("float32", ndim=1))
|
|
x2 = relax.Var("x", relax.TensorType(s0, "float32"))
|
|
x3 = relax.Var("x", relax.TensorType(s1, "float32"))
|
|
|
|
with pytest.raises(ValueError):
|
|
bb.normalize(relax.op.concat([x0], axis=1))
|
|
with pytest.raises(ValueError):
|
|
bb.normalize(relax.op.concat([x1], axis=1))
|
|
with pytest.raises(ValueError):
|
|
bb.normalize(relax.op.concat([x2], axis=1))
|
|
with pytest.raises(ValueError):
|
|
bb.normalize(relax.op.concat([x3], axis=1))
|
|
|
|
|
|
def test_concat_infer_ty_unequal_shape():
|
|
bb = relax.BlockBuilder()
|
|
a = tirx.Var("a", "int64")
|
|
s0 = relax.Var("s", relax.ShapeType((3, 4)))
|
|
s1 = relax.Var("s", relax.ShapeType((3, a + 2)))
|
|
x0 = relax.Var("x", R.Tensor((3, 4), "float32"))
|
|
x1 = relax.Var("x", R.Tensor((3, a + 2), "float32"))
|
|
x2 = relax.Var("x", relax.TensorType(s0, "float32"))
|
|
x3 = relax.Var("x", relax.TensorType(s1, "float32"))
|
|
y0 = relax.Var("y", R.Tensor((3, 3), "float32"))
|
|
y1 = relax.Var("y", R.Tensor((3, a), "float32"))
|
|
|
|
with pytest.raises(ValueError):
|
|
bb.normalize(relax.op.concat([x0, y0]))
|
|
with pytest.raises(ValueError):
|
|
bb.normalize(relax.op.concat([x2, y0]))
|
|
with pytest.raises(ValueError):
|
|
bb.normalize(relax.op.concat([x1, y1]))
|
|
with pytest.raises(ValueError):
|
|
bb.normalize(relax.op.concat([x3, y1]))
|
|
|
|
|
|
def test_concat_infer_ty_input_not_tuple():
|
|
bb = relax.BlockBuilder()
|
|
x = relax.Var("x", R.Tensor((3,), "float32"))
|
|
s = relax.Var("s", relax.ShapeType((3,)))
|
|
|
|
with pytest.raises(TypeError):
|
|
bb.normalize(relax.op.concat(x))
|
|
with pytest.raises(TypeError):
|
|
bb.normalize(relax.op.concat(s))
|
|
|
|
|
|
def test_concat_infer_ty_input_tuple_field_not_tensor():
|
|
bb = relax.BlockBuilder()
|
|
s = relax.Var("s", relax.ShapeType((3,)))
|
|
|
|
with pytest.raises(TypeError):
|
|
bb.normalize(relax.op.concat([s]))
|
|
|
|
|
|
def test_split_infer_ty_by_indices():
|
|
bb = relax.BlockBuilder()
|
|
vdev0 = VDevice("llvm")
|
|
x0 = relax.Var("x", R.Tensor((2, 10, 4), "float32"))
|
|
x1 = relax.Var("x", R.Tensor("float32", ndim=3))
|
|
x2 = relax.Var("x", R.Tensor("float32"))
|
|
x3 = relax.Var("x", R.Tensor((2, 10, 4)))
|
|
x4 = relax.Var("x", R.Tensor(ndim=3))
|
|
x5 = relax.Var("x", R.Tensor())
|
|
x6 = relax.Var("x", R.Tensor((2, 10, 4), "float32", vdev0))
|
|
|
|
_check_inference(
|
|
bb,
|
|
relax.op.split(x0, [3, 7], axis=1),
|
|
relax.TupleType(
|
|
[
|
|
relax.TensorType((2, 3, 4), "float32"),
|
|
relax.TensorType((2, 4, 4), "float32"),
|
|
relax.TensorType((2, 3, 4), "float32"),
|
|
]
|
|
),
|
|
)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.split(x6, [3, 7], axis=1),
|
|
relax.TupleType(
|
|
[
|
|
relax.TensorType((2, 3, 4), "float32", vdev0),
|
|
relax.TensorType((2, 4, 4), "float32", vdev0),
|
|
relax.TensorType((2, 3, 4), "float32", vdev0),
|
|
]
|
|
),
|
|
)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.split(x0, [3, 7], axis=-2),
|
|
relax.TupleType(
|
|
[
|
|
relax.TensorType((2, 3, 4), "float32"),
|
|
relax.TensorType((2, 4, 4), "float32"),
|
|
relax.TensorType((2, 3, 4), "float32"),
|
|
]
|
|
),
|
|
)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.split(x1, [3, 7], axis=1),
|
|
relax.TupleType(
|
|
[
|
|
relax.TensorType(dtype="float32", ndim=3),
|
|
relax.TensorType(dtype="float32", ndim=3),
|
|
relax.TensorType(dtype="float32", ndim=3),
|
|
]
|
|
),
|
|
)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.split(x2, [3, 7], axis=1),
|
|
relax.TupleType(
|
|
[
|
|
relax.TensorType(dtype="float32"),
|
|
relax.TensorType(dtype="float32"),
|
|
relax.TensorType(dtype="float32"),
|
|
]
|
|
),
|
|
)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.split(x3, [3, 7], axis=1),
|
|
relax.TupleType(
|
|
[
|
|
relax.TensorType((2, 3, 4), dtype=""),
|
|
relax.TensorType((2, 4, 4), dtype=""),
|
|
relax.TensorType((2, 3, 4), dtype=""),
|
|
]
|
|
),
|
|
)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.split(x4, [3, 7], axis=1),
|
|
relax.TupleType(
|
|
[
|
|
relax.TensorType(dtype="", ndim=3),
|
|
relax.TensorType(dtype="", ndim=3),
|
|
relax.TensorType(dtype="", ndim=3),
|
|
]
|
|
),
|
|
)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.split(x5, [3, 7], axis=1),
|
|
relax.TupleType(
|
|
[
|
|
relax.TensorType(dtype=""),
|
|
relax.TensorType(dtype=""),
|
|
relax.TensorType(dtype=""),
|
|
]
|
|
),
|
|
)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.split(x0, [-2, 2, 6, 4, 8, 12, 9], axis=1),
|
|
relax.TupleType(
|
|
[
|
|
relax.TensorType((2, 0, 4), "float32"),
|
|
relax.TensorType((2, 2, 4), "float32"),
|
|
relax.TensorType((2, 4, 4), "float32"),
|
|
relax.TensorType((2, 0, 4), "float32"),
|
|
relax.TensorType((2, 4, 4), "float32"),
|
|
relax.TensorType((2, 2, 4), "float32"),
|
|
relax.TensorType((2, 0, 4), "float32"),
|
|
relax.TensorType((2, 1, 4), "float32"),
|
|
]
|
|
),
|
|
)
|
|
|
|
|
|
def test_split_infer_ty_by_indices_shape_symbolic():
|
|
bb = relax.BlockBuilder()
|
|
a = tirx.Var("a", "int64")
|
|
b = tirx.Var("b", "int64")
|
|
x = relax.Var("x", R.Tensor((a, b), "float32"))
|
|
|
|
_check_inference(
|
|
bb,
|
|
relax.op.split(x, [10, 20], axis=1),
|
|
relax.TupleType(
|
|
[
|
|
relax.TensorType([a, T.max(T.min(10, b) - T.min(0, b), 0)], dtype="float32"),
|
|
relax.TensorType([a, T.max(T.min(20, b) - T.min(10, b), 0)], dtype="float32"),
|
|
relax.TensorType([a, T.max(b - 20, 0)], dtype="float32"),
|
|
]
|
|
),
|
|
)
|
|
|
|
|
|
def test_split_infer_ty_by_indices_shape_var():
|
|
bb = relax.BlockBuilder()
|
|
s0 = relax.Var("s", relax.ShapeType((2, 10, 4)))
|
|
s1 = relax.Var("s", relax.ShapeType(ndim=3))
|
|
s2 = relax.Var("s", relax.ShapeType())
|
|
x0 = relax.Var("x", relax.TensorType(s0, "float32"))
|
|
x1 = relax.Var("x", relax.TensorType(s1, "float32"))
|
|
x2 = relax.Var("x", relax.TensorType(s2, "float32"))
|
|
|
|
_check_inference(
|
|
bb,
|
|
relax.op.split(x0, [3], axis=1),
|
|
relax.TupleType(
|
|
[
|
|
relax.TensorType(dtype="float32", ndim=3),
|
|
relax.TensorType(dtype="float32", ndim=3),
|
|
]
|
|
),
|
|
)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.split(x1, [3], axis=1),
|
|
relax.TupleType(
|
|
[
|
|
relax.TensorType(dtype="float32", ndim=3),
|
|
relax.TensorType(dtype="float32", ndim=3),
|
|
]
|
|
),
|
|
)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.split(x2, [3], axis=1),
|
|
relax.TupleType([relax.TensorType(dtype="float32"), relax.TensorType(dtype="float32")]),
|
|
)
|
|
|
|
|
|
def test_split_infer_ty_by_n_section():
|
|
bb = relax.BlockBuilder()
|
|
x0 = relax.Var("x", R.Tensor((2, 10, 4), "float32"))
|
|
x1 = relax.Var("x", R.Tensor("float32", ndim=3))
|
|
x2 = relax.Var("x", R.Tensor("float32"))
|
|
x3 = relax.Var("x", R.Tensor((2, 10, 4)))
|
|
x4 = relax.Var("x", R.Tensor(ndim=3))
|
|
x5 = relax.Var("x", R.Tensor())
|
|
|
|
_check_inference(
|
|
bb,
|
|
relax.op.split(x0, 3, axis=1),
|
|
relax.TupleType(
|
|
[
|
|
relax.TensorType((2, 4, 4), "float32"),
|
|
relax.TensorType((2, 4, 4), "float32"),
|
|
relax.TensorType((2, 2, 4), "float32"),
|
|
]
|
|
),
|
|
)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.split(x0, 2, axis=1),
|
|
relax.TupleType(
|
|
[
|
|
relax.TensorType((2, 5, 4), "float32"),
|
|
relax.TensorType((2, 5, 4), "float32"),
|
|
]
|
|
),
|
|
)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.split(x0, 3, axis=-2),
|
|
relax.TupleType(
|
|
[
|
|
relax.TensorType((2, 4, 4), "float32"),
|
|
relax.TensorType((2, 4, 4), "float32"),
|
|
relax.TensorType((2, 2, 4), "float32"),
|
|
]
|
|
),
|
|
)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.split(x1, 3, axis=1),
|
|
relax.TupleType(
|
|
[
|
|
relax.TensorType(dtype="float32", ndim=3),
|
|
relax.TensorType(dtype="float32", ndim=3),
|
|
relax.TensorType(dtype="float32", ndim=3),
|
|
]
|
|
),
|
|
)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.split(x2, 3, axis=1),
|
|
relax.TupleType(
|
|
[
|
|
relax.TensorType(dtype="float32"),
|
|
relax.TensorType(dtype="float32"),
|
|
relax.TensorType(dtype="float32"),
|
|
]
|
|
),
|
|
)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.split(x3, 3, axis=1),
|
|
relax.TupleType(
|
|
[
|
|
relax.TensorType((2, 4, 4), dtype=""),
|
|
relax.TensorType((2, 4, 4), dtype=""),
|
|
relax.TensorType((2, 2, 4), dtype=""),
|
|
]
|
|
),
|
|
)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.split(x4, 3, axis=1),
|
|
relax.TupleType(
|
|
[
|
|
relax.TensorType(dtype="", ndim=3),
|
|
relax.TensorType(dtype="", ndim=3),
|
|
relax.TensorType(dtype="", ndim=3),
|
|
]
|
|
),
|
|
)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.split(x5, 3, axis=1),
|
|
relax.TupleType(
|
|
[
|
|
relax.TensorType(dtype=""),
|
|
relax.TensorType(dtype=""),
|
|
relax.TensorType(dtype=""),
|
|
]
|
|
),
|
|
)
|
|
|
|
|
|
def test_split_infer_ty_by_n_section_shape_symbolic():
|
|
bb = relax.BlockBuilder()
|
|
a = tirx.Var("a", "int64")
|
|
b = tirx.Var("b", "int64")
|
|
x = relax.Var("x", R.Tensor((a, b), "float32"))
|
|
|
|
_check_inference(
|
|
bb,
|
|
relax.op.split(x, 3, axis=1),
|
|
relax.TupleType(
|
|
[
|
|
relax.TensorType((a, (b + 2) // 3), "float32"),
|
|
relax.TensorType((a, (b + 2) // 3), "float32"),
|
|
relax.TensorType((a, b - (b + 2) // 3 * 2), "float32"),
|
|
]
|
|
),
|
|
)
|
|
|
|
|
|
def test_split_infer_ty_by_n_section_shape_var():
|
|
bb = relax.BlockBuilder()
|
|
s0 = relax.Var("s", relax.ShapeType((2, 10, 4)))
|
|
s1 = relax.Var("s", relax.ShapeType(ndim=3))
|
|
s2 = relax.Var("s", relax.ShapeType())
|
|
x0 = relax.Var("x", relax.TensorType(s0, "float32"))
|
|
x1 = relax.Var("x", relax.TensorType(s1, "float32"))
|
|
x2 = relax.Var("x", relax.TensorType(s2, "float32"))
|
|
|
|
_check_inference(
|
|
bb,
|
|
relax.op.split(x0, 3, axis=1),
|
|
relax.TupleType(
|
|
[
|
|
relax.TensorType(dtype="float32", ndim=3),
|
|
relax.TensorType(dtype="float32", ndim=3),
|
|
relax.TensorType(dtype="float32", ndim=3),
|
|
]
|
|
),
|
|
)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.split(x1, 3, axis=1),
|
|
relax.TupleType(
|
|
[
|
|
relax.TensorType(dtype="float32", ndim=3),
|
|
relax.TensorType(dtype="float32", ndim=3),
|
|
relax.TensorType(dtype="float32", ndim=3),
|
|
]
|
|
),
|
|
)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.split(x2, 3, axis=1),
|
|
relax.TupleType(
|
|
[
|
|
relax.TensorType(dtype="float32"),
|
|
relax.TensorType(dtype="float32"),
|
|
relax.TensorType(dtype="float32"),
|
|
]
|
|
),
|
|
)
|
|
|
|
|
|
def test_split_infer_ty_more_input_dtype():
|
|
bb = relax.BlockBuilder()
|
|
x0 = relax.Var("x", R.Tensor((2, 10, 4), "float16"))
|
|
x1 = relax.Var("x", R.Tensor((2, 10, 4), "int8"))
|
|
|
|
_check_inference(
|
|
bb,
|
|
relax.op.split(x0, [3, 7], axis=1),
|
|
relax.TupleType(
|
|
[
|
|
relax.TensorType((2, 3, 4), "float16"),
|
|
relax.TensorType((2, 4, 4), "float16"),
|
|
relax.TensorType((2, 3, 4), "float16"),
|
|
]
|
|
),
|
|
)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.split(x1, [3, 7], axis=1),
|
|
relax.TupleType(
|
|
[
|
|
relax.TensorType((2, 3, 4), "int8"),
|
|
relax.TensorType((2, 4, 4), "int8"),
|
|
relax.TensorType((2, 3, 4), "int8"),
|
|
]
|
|
),
|
|
)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.split(x0, 3, axis=1),
|
|
relax.TupleType(
|
|
[
|
|
relax.TensorType((2, 4, 4), "float16"),
|
|
relax.TensorType((2, 4, 4), "float16"),
|
|
relax.TensorType((2, 2, 4), "float16"),
|
|
]
|
|
),
|
|
)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.split(x1, 3, axis=1),
|
|
relax.TupleType(
|
|
[
|
|
relax.TensorType((2, 4, 4), "int8"),
|
|
relax.TensorType((2, 4, 4), "int8"),
|
|
relax.TensorType((2, 2, 4), "int8"),
|
|
]
|
|
),
|
|
)
|
|
|
|
|
|
def test_split_infer_ty_single_output():
|
|
bb = relax.BlockBuilder()
|
|
a = tirx.Var("a", "int64")
|
|
b = tirx.Var("b", "int64")
|
|
s0 = relax.Var("s", relax.ShapeType((a, b)))
|
|
s1 = relax.Var("s", relax.ShapeType(ndim=2))
|
|
s2 = relax.Var("s", relax.ShapeType())
|
|
x0 = relax.Var("x", R.Tensor((a, b), "float32"))
|
|
x1 = relax.Var("x", R.Tensor("float32", ndim=2))
|
|
x2 = relax.Var("x", R.Tensor("float32"))
|
|
x3 = relax.Var("x", relax.TensorType(s0, "float32"))
|
|
x4 = relax.Var("x", relax.TensorType(s1, "float32"))
|
|
x5 = relax.Var("x", relax.TensorType(s2, "float32"))
|
|
|
|
_check_inference(
|
|
bb,
|
|
relax.op.split(x0, [], axis=1),
|
|
relax.TensorType((a, b), "float32"),
|
|
)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.split(x1, [], axis=1),
|
|
relax.TensorType(dtype="float32", ndim=2),
|
|
)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.split(x2, [], axis=1),
|
|
relax.TensorType(dtype="float32"),
|
|
)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.split(x3, [], axis=1),
|
|
relax.TensorType(s0, "float32"),
|
|
)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.split(x4, [], axis=1),
|
|
relax.TensorType(s1, "float32"),
|
|
)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.split(x5, [], axis=1),
|
|
relax.TensorType(s2, "float32"),
|
|
)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.split(x0, 1, axis=1),
|
|
relax.TensorType((a, b), "float32"),
|
|
)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.split(x1, 1, axis=1),
|
|
relax.TensorType(dtype="float32", ndim=2),
|
|
)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.split(x2, 1, axis=1),
|
|
relax.TensorType(dtype="float32"),
|
|
)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.split(x3, 1, axis=1),
|
|
relax.TensorType(s0, "float32"),
|
|
)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.split(x4, 1, axis=1),
|
|
relax.TensorType(s1, "float32"),
|
|
)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.split(x5, 1, axis=1),
|
|
relax.TensorType(s2, "float32"),
|
|
)
|
|
|
|
|
|
def test_split_indices_or_sections_int64():
|
|
x = relax.Var("x", R.Tensor((2, 10, 4), "float32"))
|
|
split0 = relax.op.split(x, [3, 6], axis=1)
|
|
split1 = relax.op.split(x, 4, axis=1)
|
|
|
|
assert split0.attrs.indices_or_sections[0].ty.dtype == "int64"
|
|
assert split0.attrs.indices_or_sections[1].ty.dtype == "int64"
|
|
assert split1.attrs.indices_or_sections.ty.dtype == "int64"
|
|
|
|
|
|
def test_split_infer_ty():
|
|
bb = relax.BlockBuilder()
|
|
n = tirx.Var("n", "int64")
|
|
x = relax.Var("x", R.Tensor((16, 4)))
|
|
y = relax.Var("y", R.Tensor((16, 4), "float32"))
|
|
z = relax.Var("z", R.Tensor((n, 16)))
|
|
w = relax.Var("w", R.Tensor((n + 5, 16)))
|
|
|
|
# All relax shape variables are non-negative. When a scope
|
|
# begins, any TIR variables that are used as shape variables are
|
|
# declared to be non-negative `tvm.arith.Analyzer`. Because
|
|
# `relax.op.split` clamps the indices to be within the bounds of
|
|
# the axis being split, simplifying with non-negative shape
|
|
# variables can result in much simpler shapes.
|
|
#
|
|
# For example, an axis of size `n`, split on the range from 2 to 5
|
|
# has size `T.max(T.min(5, n + 5) - T.min(2, n + 5), 0)`. If it
|
|
# is known that `n >= 0`, then this simplifies down to `3`.
|
|
bb.begin_scope([x, y, z, w])
|
|
|
|
_check_inference(
|
|
bb,
|
|
relax.op.split(x, 1),
|
|
R.Tensor([16, 4]),
|
|
)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.split(x, 2),
|
|
R.Tuple(
|
|
R.Tensor([8, 4]),
|
|
R.Tensor([8, 4]),
|
|
),
|
|
)
|
|
# Uneven splits are allowed, with the last split being smaller than the others.
|
|
_check_inference(
|
|
bb,
|
|
relax.op.split(x, 3),
|
|
R.Tuple(
|
|
R.Tensor([6, 4]),
|
|
R.Tensor([6, 4]),
|
|
R.Tensor([4, 4]),
|
|
),
|
|
)
|
|
|
|
# Dtype of result is inherited from the tensor
|
|
_check_inference(
|
|
bb,
|
|
relax.op.split(y, 2),
|
|
R.Tuple(
|
|
R.Tensor([8, 4], "float32"),
|
|
R.Tensor([8, 4], "float32"),
|
|
),
|
|
)
|
|
|
|
# Axis can be explicitly specified. Otherwise, defaults to axis=0.
|
|
_check_inference(
|
|
bb, relax.op.split(x, [2], axis=1), R.Tuple(R.Tensor([16, 2]), R.Tensor([16, 2]))
|
|
)
|
|
|
|
# Split points can be explicitly specified
|
|
_check_inference(
|
|
bb,
|
|
relax.op.split(x, [2]),
|
|
R.Tuple(
|
|
R.Tensor([2, 4]),
|
|
R.Tensor([14, 4]),
|
|
),
|
|
)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.split(x, [2, 5]),
|
|
R.Tuple(
|
|
R.Tensor([2, 4]),
|
|
R.Tensor([3, 4]),
|
|
R.Tensor([11, 4]),
|
|
),
|
|
)
|
|
|
|
# Splitting a dynamic axis is allowed, and propagates the shape to the output
|
|
_check_inference(
|
|
bb,
|
|
relax.op.split(z, 2),
|
|
R.Tuple(
|
|
R.Tensor([(n + 1) // 2, 16]),
|
|
R.Tensor([n - (n + 1) // 2, 16]),
|
|
),
|
|
)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.split(z, 3),
|
|
R.Tuple(
|
|
R.Tensor([(n + 2) // 3, 16]),
|
|
R.Tensor([(n + 2) // 3, 16]),
|
|
R.Tensor([n - (n + 2) // 3 * 2, 16]),
|
|
),
|
|
)
|
|
|
|
# Splitting a dynamic axis at specific indices is allowed.
|
|
_check_inference(
|
|
bb,
|
|
relax.op.split(w, [2, 5]),
|
|
R.Tuple(
|
|
R.Tensor((2, 16)),
|
|
R.Tensor((3, 16)),
|
|
R.Tensor((n, 16)),
|
|
),
|
|
)
|
|
|
|
|
|
def test_split_infer_ty_non_integer_indices():
|
|
bb = relax.BlockBuilder()
|
|
a = tirx.Var("c", "int64")
|
|
b = tirx.Var("d", "int64")
|
|
x = relax.Var("x", R.Tensor((3, 4), "float32"))
|
|
|
|
with pytest.raises(TypeError):
|
|
bb.normalize(relax.op.split(x, [a, b], axis=1))
|
|
|
|
|
|
def test_split_invalid_n_section():
|
|
n = tirx.Var("n", "int64")
|
|
x = relax.Var("x", R.Tensor((3, 4), "float32"))
|
|
|
|
with pytest.raises(tvm.error.InternalError):
|
|
relax.op.split(x, 0, axis=1)
|
|
with pytest.raises(tvm.error.InternalError):
|
|
relax.op.split(x, -1, axis=1)
|
|
with pytest.raises(TypeError):
|
|
relax.op.split(x, n, axis=1)
|
|
|
|
|
|
def test_split_infer_ty_axis_out_of_range():
|
|
bb = relax.BlockBuilder()
|
|
x0 = relax.Var("x", R.Tensor((2, 3), "float32"))
|
|
x1 = relax.Var("x", R.Tensor("float32", ndim=2))
|
|
|
|
with pytest.raises(ValueError):
|
|
bb.normalize(relax.op.split(x0, [], axis=2))
|
|
with pytest.raises(ValueError):
|
|
bb.normalize(relax.op.split(x0, [], axis=-3))
|
|
with pytest.raises(ValueError):
|
|
bb.normalize(relax.op.split(x1, 1, axis=2))
|
|
with pytest.raises(ValueError):
|
|
bb.normalize(relax.op.split(x1, 1, axis=-3))
|
|
|
|
|
|
def test_split_infer_invalid_ty_indices():
|
|
bb = relax.BlockBuilder()
|
|
x0 = relax.Var("x", R.Tensor((2, 3), "float32"))
|
|
v = relax.Var("v", tvm.ir.PrimType("int64"))
|
|
|
|
with pytest.raises(TypeError):
|
|
bb.normalize(relax.op.split(x0, [v], axis=1))
|
|
with pytest.raises(TypeError):
|
|
bb.normalize(relax.op.split(x0, v, axis=1))
|
|
|
|
|
|
def test_split_infer_ty_wrong_input_type():
|
|
bb = relax.BlockBuilder()
|
|
x0 = relax.Var("x", relax.ShapeType((2, 3)))
|
|
x1 = relax.Var("x", relax.FuncType([], R.Tensor((2, 3), "float32")))
|
|
|
|
with pytest.raises(TypeError):
|
|
bb.normalize(relax.op.split(x0, 1, axis=1))
|
|
with pytest.raises(TypeError):
|
|
bb.normalize(relax.op.split(x1, 1, axis=1))
|
|
|
|
|
|
def test_broadcast_to_infer_ty():
|
|
bb = relax.BlockBuilder()
|
|
vdev0 = VDevice("llvm")
|
|
x0 = relax.Var("x", R.Tensor((2, 1, 3), "float32"))
|
|
x1 = relax.Var("x", R.Tensor("float32", ndim=3))
|
|
x2 = relax.Var("x", R.Tensor("float32"))
|
|
x3 = relax.Var("x", R.Tensor((2, 1, 3)))
|
|
x4 = relax.Var("x", R.Tensor(ndim=3))
|
|
x5 = relax.Var("x", R.Tensor())
|
|
x6 = relax.Var("x", R.Tensor((2, 1, 3), "float32", vdev0))
|
|
|
|
_check_inference(
|
|
bb, relax.op.broadcast_to(x0, (4, 2, 5, 3)), relax.TensorType((4, 2, 5, 3), "float32")
|
|
)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.broadcast_to(x6, (4, 2, 5, 3)),
|
|
relax.TensorType((4, 2, 5, 3), "float32", vdev0),
|
|
)
|
|
_check_inference(
|
|
bb, relax.op.broadcast_to(x1, (4, 2, 5, 3)), relax.TensorType((4, 2, 5, 3), "float32")
|
|
)
|
|
_check_inference(
|
|
bb, relax.op.broadcast_to(x2, (4, 2, 5, 3)), relax.TensorType((4, 2, 5, 3), "float32")
|
|
)
|
|
_check_inference(
|
|
bb, relax.op.broadcast_to(x3, (4, 2, 5, 3)), relax.TensorType((4, 2, 5, 3), dtype="")
|
|
)
|
|
_check_inference(
|
|
bb, relax.op.broadcast_to(x4, (4, 2, 5, 3)), relax.TensorType((4, 2, 5, 3), dtype="")
|
|
)
|
|
_check_inference(
|
|
bb, relax.op.broadcast_to(x5, (4, 2, 5, 3)), relax.TensorType((4, 2, 5, 3), dtype="")
|
|
)
|
|
|
|
|
|
def test_broadcast_to_infer_ty_shape_symbolic():
|
|
bb = relax.BlockBuilder()
|
|
a = tirx.Var("a", "int64")
|
|
b = tirx.Var("b", "int64")
|
|
c = tirx.Var("c", "int64")
|
|
d = tirx.Var("d", "int64")
|
|
x0 = relax.Var("x", R.Tensor((b, 1, 1, d), "float32"))
|
|
x1 = relax.Var("x", R.Tensor((b, 1, 1, d)))
|
|
|
|
_check_inference(
|
|
bb,
|
|
relax.op.broadcast_to(x0, (a, b, 1, c, d)),
|
|
relax.TensorType((a, b, 1, c, d), "float32"),
|
|
)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.broadcast_to(x1, (a, b, 1, c, d)),
|
|
relax.TensorType((a, b, 1, c, d), dtype=""),
|
|
)
|
|
|
|
|
|
def test_broadcast_to_infer_ty_shape_var():
|
|
bb = relax.BlockBuilder()
|
|
s0 = relax.Var("s", relax.ShapeType((2, 1, 3)))
|
|
s1 = relax.Var("s", relax.ShapeType(ndim=3))
|
|
s2 = relax.Var("s", relax.ShapeType())
|
|
x0 = relax.Var("x", relax.TensorType(s0, "float32"))
|
|
x1 = relax.Var("x", relax.TensorType(s1, "float32"))
|
|
x2 = relax.Var("x", relax.TensorType(s2, "float32"))
|
|
|
|
_check_inference(
|
|
bb, relax.op.broadcast_to(x0, (4, 2, 5, 3)), relax.TensorType((4, 2, 5, 3), "float32")
|
|
)
|
|
_check_inference(
|
|
bb, relax.op.broadcast_to(x1, (4, 2, 5, 3)), relax.TensorType((4, 2, 5, 3), "float32")
|
|
)
|
|
_check_inference(
|
|
bb, relax.op.broadcast_to(x2, (4, 2, 5, 3)), relax.TensorType((4, 2, 5, 3), "float32")
|
|
)
|
|
|
|
|
|
def test_broadcast_to_infer_ty_tgt_shape_var():
|
|
bb = relax.BlockBuilder()
|
|
a = tirx.Var("a", "int64")
|
|
b = tirx.Var("b", "int64")
|
|
c = tirx.Var("c", "int64")
|
|
d = tirx.Var("d", "int64")
|
|
s0 = relax.Var("s", relax.ShapeType((b, 1, 1, d)))
|
|
s1 = relax.Var("s", relax.ShapeType(ndim=4))
|
|
s2 = relax.Var("s", relax.ShapeType())
|
|
x0 = relax.Var("x", R.Tensor((b, 1, 1, d), "float32"))
|
|
x1 = relax.Var("x", R.Tensor("float32", ndim=4))
|
|
x2 = relax.Var("x", R.Tensor("float32"))
|
|
x3 = relax.Var("x", relax.TensorType(s0, "float32"))
|
|
x4 = relax.Var("x", relax.TensorType(s1, "float32"))
|
|
x5 = relax.Var("x", relax.TensorType(s2, "float32"))
|
|
stgt0 = relax.Var("stgt", relax.ShapeType((a, b, 1, c, d)))
|
|
stgt1 = relax.Var("stgt", relax.ShapeType(ndim=5))
|
|
stgt2 = relax.Var("stgt", relax.ShapeType())
|
|
|
|
_check_inference(bb, relax.op.broadcast_to(x0, stgt0), relax.TensorType(stgt0, "float32"))
|
|
_check_inference(bb, relax.op.broadcast_to(x1, stgt0), relax.TensorType(stgt0, "float32"))
|
|
_check_inference(bb, relax.op.broadcast_to(x2, stgt0), relax.TensorType(stgt0, "float32"))
|
|
_check_inference(bb, relax.op.broadcast_to(x3, stgt0), relax.TensorType(stgt0, "float32"))
|
|
_check_inference(bb, relax.op.broadcast_to(x4, stgt0), relax.TensorType(stgt0, "float32"))
|
|
_check_inference(bb, relax.op.broadcast_to(x5, stgt0), relax.TensorType(stgt0, "float32"))
|
|
_check_inference(bb, relax.op.broadcast_to(x0, stgt1), relax.TensorType(stgt1, "float32"))
|
|
_check_inference(bb, relax.op.broadcast_to(x1, stgt1), relax.TensorType(stgt1, "float32"))
|
|
_check_inference(bb, relax.op.broadcast_to(x2, stgt1), relax.TensorType(stgt1, "float32"))
|
|
_check_inference(bb, relax.op.broadcast_to(x3, stgt1), relax.TensorType(stgt1, "float32"))
|
|
_check_inference(bb, relax.op.broadcast_to(x4, stgt1), relax.TensorType(stgt1, "float32"))
|
|
_check_inference(bb, relax.op.broadcast_to(x5, stgt1), relax.TensorType(stgt1, "float32"))
|
|
_check_inference(bb, relax.op.broadcast_to(x0, stgt2), relax.TensorType(stgt2, "float32"))
|
|
_check_inference(bb, relax.op.broadcast_to(x1, stgt2), relax.TensorType(stgt2, "float32"))
|
|
_check_inference(bb, relax.op.broadcast_to(x2, stgt2), relax.TensorType(stgt2, "float32"))
|
|
_check_inference(bb, relax.op.broadcast_to(x3, stgt2), relax.TensorType(stgt2, "float32"))
|
|
_check_inference(bb, relax.op.broadcast_to(x4, stgt2), relax.TensorType(stgt2, "float32"))
|
|
_check_inference(bb, relax.op.broadcast_to(x5, stgt2), relax.TensorType(stgt2, "float32"))
|
|
|
|
|
|
def test_broadcast_to_infer_ty_more_input_dtype():
|
|
bb = relax.BlockBuilder()
|
|
x0 = relax.Var("x", R.Tensor((2, 1, 3), "float16"))
|
|
x1 = relax.Var("x", R.Tensor((2, 1, 3), "int8"))
|
|
x2 = relax.Var("x", R.Tensor((2, 1, 3), "int32"))
|
|
|
|
_check_inference(
|
|
bb, relax.op.broadcast_to(x0, (4, 2, 5, 3)), relax.TensorType((4, 2, 5, 3), "float16")
|
|
)
|
|
_check_inference(
|
|
bb, relax.op.broadcast_to(x1, (4, 2, 5, 3)), relax.TensorType((4, 2, 5, 3), "int8")
|
|
)
|
|
_check_inference(
|
|
bb, relax.op.broadcast_to(x2, (4, 2, 5, 3)), relax.TensorType((4, 2, 5, 3), "int32")
|
|
)
|
|
|
|
|
|
def test_broadcast_to_infer_ty_tgt_ndim_less_than_old_ndim():
|
|
bb = relax.BlockBuilder()
|
|
s0 = relax.Var("s", relax.ShapeType((2, 1)))
|
|
s1 = relax.Var("s", relax.ShapeType(ndim=2))
|
|
x0 = relax.Var("x", R.Tensor((2, 1), "float32"))
|
|
x1 = relax.Var("x", R.Tensor("float32", ndim=2))
|
|
x2 = relax.Var("x", relax.TensorType(s0, "float32"))
|
|
x3 = relax.Var("x", relax.TensorType(s1, "float32"))
|
|
stgt0 = relax.Var("stgt", relax.ShapeType((2,)))
|
|
stgt1 = relax.Var("stgt", relax.ShapeType(ndim=1))
|
|
|
|
with pytest.raises(ValueError):
|
|
bb.normalize(relax.op.broadcast_to(x0, (2,)))
|
|
with pytest.raises(ValueError):
|
|
bb.normalize(relax.op.broadcast_to(x0, stgt0))
|
|
with pytest.raises(ValueError):
|
|
bb.normalize(relax.op.broadcast_to(x0, stgt1))
|
|
with pytest.raises(ValueError):
|
|
bb.normalize(relax.op.broadcast_to(x1, (2,)))
|
|
with pytest.raises(ValueError):
|
|
bb.normalize(relax.op.broadcast_to(x1, stgt0))
|
|
with pytest.raises(ValueError):
|
|
bb.normalize(relax.op.broadcast_to(x1, stgt1))
|
|
with pytest.raises(ValueError):
|
|
bb.normalize(relax.op.broadcast_to(x2, (2,)))
|
|
with pytest.raises(ValueError):
|
|
bb.normalize(relax.op.broadcast_to(x2, stgt0))
|
|
with pytest.raises(ValueError):
|
|
bb.normalize(relax.op.broadcast_to(x2, stgt1))
|
|
with pytest.raises(ValueError):
|
|
bb.normalize(relax.op.broadcast_to(x3, (2,)))
|
|
with pytest.raises(ValueError):
|
|
bb.normalize(relax.op.broadcast_to(x3, stgt0))
|
|
with pytest.raises(ValueError):
|
|
bb.normalize(relax.op.broadcast_to(x3, stgt1))
|
|
|
|
|
|
def test_broadcast_to_infer_ty_not_broadcastable_static():
|
|
bb = relax.BlockBuilder()
|
|
s = relax.Var("s", relax.ShapeType((2, 1, 3)))
|
|
x0 = relax.Var("x", R.Tensor((2, 1, 3), "float32"))
|
|
x1 = relax.Var("x", relax.TensorType(s, "float32"))
|
|
stgt = relax.Var("stgt", relax.ShapeType((2, 1, 6)))
|
|
|
|
with pytest.raises(ValueError):
|
|
bb.normalize(relax.op.broadcast_to(x0, (2, 1, 6)))
|
|
with pytest.raises(ValueError):
|
|
bb.normalize(relax.op.broadcast_to(x0, stgt))
|
|
with pytest.raises(ValueError):
|
|
bb.normalize(relax.op.broadcast_to(x1, (2, 1, 6)))
|
|
with pytest.raises(ValueError):
|
|
bb.normalize(relax.op.broadcast_to(x1, stgt))
|
|
|
|
|
|
def test_broadcast_to_infer_ty_not_broadcastable_symbolic():
|
|
bb = relax.BlockBuilder()
|
|
a = tirx.Var("a", "int64")
|
|
b = tirx.Var("b", "int64")
|
|
s = relax.Var("s", relax.ShapeType((2, a)))
|
|
x0 = relax.Var("x", R.Tensor((2, a), "float32"))
|
|
x1 = relax.Var("x", relax.TensorType(s, "float32"))
|
|
stgt0 = relax.Var("stgt", relax.ShapeType((2, b)))
|
|
stgt1 = relax.Var("stgt", relax.ShapeType((2, 1)))
|
|
stgt2 = relax.Var("stgt", relax.ShapeType((b, a)))
|
|
|
|
_check_inference(bb, relax.op.broadcast_to(x0, (2, b)), relax.TensorType((2, b), "float32"))
|
|
_check_inference(bb, relax.op.broadcast_to(x0, (2, 1)), relax.TensorType((2, 1), "float32"))
|
|
_check_inference(bb, relax.op.broadcast_to(x0, (b, a)), relax.TensorType((b, a), "float32"))
|
|
_check_inference(bb, relax.op.broadcast_to(x0, stgt0), relax.TensorType(stgt0, "float32"))
|
|
_check_inference(bb, relax.op.broadcast_to(x0, stgt1), relax.TensorType(stgt1, "float32"))
|
|
_check_inference(bb, relax.op.broadcast_to(x0, stgt2), relax.TensorType(stgt2, "float32"))
|
|
_check_inference(bb, relax.op.broadcast_to(x1, (2, b)), relax.TensorType((2, b), "float32"))
|
|
_check_inference(bb, relax.op.broadcast_to(x1, (2, 1)), relax.TensorType((2, 1), "float32"))
|
|
_check_inference(bb, relax.op.broadcast_to(x1, (b, a)), relax.TensorType((b, a), "float32"))
|
|
_check_inference(bb, relax.op.broadcast_to(x1, stgt0), relax.TensorType(stgt0, "float32"))
|
|
_check_inference(bb, relax.op.broadcast_to(x1, stgt1), relax.TensorType(stgt1, "float32"))
|
|
_check_inference(bb, relax.op.broadcast_to(x1, stgt2), relax.TensorType(stgt2, "float32"))
|
|
|
|
|
|
def test_broadcast_to_infer_ty_wrong_input_type():
|
|
bb = relax.BlockBuilder()
|
|
x0 = relax.Var("x", relax.ShapeType((2, 1, 3)))
|
|
x1 = relax.Var("x", R.Tensor((2, 1, 3), "float32"))
|
|
stgt = relax.Var("stgt", relax.TensorType((4, 2, 5, 3), dtype=""))
|
|
|
|
with pytest.raises(TypeError):
|
|
bb.normalize(relax.op.broadcast_to(x0, (4, 2, 5, 3)))
|
|
with pytest.raises(TypeError):
|
|
bb.normalize(relax.op.broadcast_to(x1, stgt))
|
|
|
|
|
|
def test_collapse_sum_like_infer_ty():
|
|
bb = relax.BlockBuilder()
|
|
vdev0 = VDevice("llvm")
|
|
x0 = relax.Var("x", R.Tensor((2, 3, 4), "float32"))
|
|
x1 = relax.Var("x", R.Tensor("float32", ndim=3))
|
|
x2 = relax.Var("x", R.Tensor("float32"))
|
|
x3 = relax.Var("x", R.Tensor((2, 3, 4), "float32", vdev0))
|
|
x4 = relax.Var("x", R.Tensor(ndim=3))
|
|
x5 = relax.Var("x", R.Tensor())
|
|
y0 = relax.Var("y", R.Tensor((3, 4), "float32"))
|
|
y1 = relax.Var("y", R.Tensor("float32", ndim=2))
|
|
y2 = relax.Var("y", R.Tensor("float32"))
|
|
y3 = relax.Var("y", R.Tensor((3, 4)))
|
|
y4 = relax.Var("y", R.Tensor(ndim=2))
|
|
y5 = relax.Var("y", R.Tensor((1, 4)))
|
|
y6 = relax.Var("y", R.Tensor((3, 4), "float32", vdev0))
|
|
|
|
_check_inference(bb, relax.op.collapse_sum_like(x0, y0), relax.TensorType((3, 4), "float32"))
|
|
_check_inference(
|
|
bb, relax.op.collapse_sum_like(x3, y6), relax.TensorType((3, 4), "float32", vdev0)
|
|
)
|
|
_check_inference(
|
|
bb, relax.op.collapse_sum_like(x1, y1), relax.TensorType(dtype="float32", ndim=2)
|
|
)
|
|
_check_inference(
|
|
bb, relax.op.collapse_sum_like(x0, y1), relax.TensorType(dtype="float32", ndim=2)
|
|
)
|
|
_check_inference(
|
|
bb, relax.op.collapse_sum_like(x0, y2), relax.TensorType(dtype="float32", ndim=-1)
|
|
)
|
|
_check_inference(bb, relax.op.collapse_sum_like(x0, y3), relax.TensorType((3, 4), "float32"))
|
|
_check_inference(bb, relax.op.collapse_sum_like(x2, y0), relax.TensorType((3, 4), "float32"))
|
|
_check_inference(
|
|
bb, relax.op.collapse_sum_like(x2, y4), relax.TensorType(dtype="float32", ndim=2)
|
|
)
|
|
_check_inference(bb, relax.op.collapse_sum_like(x4, y1), relax.TensorType(dtype="", ndim=2))
|
|
_check_inference(bb, relax.op.collapse_sum_like(x5, y3), relax.TensorType((3, 4), dtype=""))
|
|
_check_inference(bb, relax.op.collapse_sum_like(x0, y5), relax.TensorType((1, 4), "float32"))
|
|
|
|
|
|
def test_collapse_sum_like_infer_ty_shape_symbolic():
|
|
bb = relax.BlockBuilder()
|
|
a = tirx.Var("a", "int64")
|
|
b = tirx.Var("b", "int64")
|
|
x0 = relax.Var("x", R.Tensor((3, 4, a), "float32"))
|
|
y0 = relax.Var("y", R.Tensor((4, a), "float32"))
|
|
x1 = relax.Var("x", R.Tensor((3, 4, b + a), "float32"))
|
|
y1 = relax.Var("x", R.Tensor((1, a + b), "float32"))
|
|
|
|
_check_inference(bb, relax.op.collapse_sum_like(x0, y0), relax.TensorType((4, a), "float32"))
|
|
_check_inference(
|
|
bb, relax.op.collapse_sum_like(x1, y1), relax.TensorType((1, a + b), "float32")
|
|
)
|
|
|
|
|
|
def test_collapse_sum_like_infer_ty_shape_var():
|
|
bb = relax.BlockBuilder()
|
|
s0 = relax.Var("s0", relax.ShapeType((2, 3, 4)))
|
|
s1 = relax.Var("s1", relax.ShapeType(ndim=3))
|
|
s2 = relax.Var("s2", relax.ShapeType())
|
|
s3 = relax.Var("s3", relax.ShapeType((3, 4)))
|
|
s4 = relax.Var("s4", relax.ShapeType(ndim=2))
|
|
s5 = relax.Var("s5", relax.ShapeType())
|
|
x0 = relax.Var("x", relax.TensorType(s0, "float32"))
|
|
x1 = relax.Var("x", relax.TensorType(s1, "float32"))
|
|
x2 = relax.Var("x", relax.TensorType(s2, "float32"))
|
|
y0 = relax.Var("y", relax.TensorType(s3, "float32"))
|
|
y1 = relax.Var("y", relax.TensorType(s4, "float32"))
|
|
y2 = relax.Var("y", relax.TensorType(s5, "float32"))
|
|
|
|
_check_inference(bb, relax.op.collapse_sum_like(x0, y0), relax.TensorType(s3, "float32"))
|
|
_check_inference(bb, relax.op.collapse_sum_like(x1, y1), relax.TensorType(s4, "float32"))
|
|
_check_inference(bb, relax.op.collapse_sum_like(x2, y2), relax.TensorType(s5, "float32"))
|
|
|
|
|
|
def test_collapse_sum_like_infer_ty_more_input_dtype():
|
|
bb = relax.BlockBuilder()
|
|
x0 = relax.Var("x", R.Tensor((2, 3, 4), "float16"))
|
|
x1 = relax.Var("x", R.Tensor((2, 3, 4), "int8"))
|
|
y0 = relax.Var("y", R.Tensor((3, 4), "float16"))
|
|
y1 = relax.Var("y", R.Tensor((3, 4), "int8"))
|
|
|
|
_check_inference(bb, relax.op.collapse_sum_like(x0, y0), relax.TensorType((3, 4), "float16"))
|
|
_check_inference(bb, relax.op.collapse_sum_like(x1, y1), relax.TensorType((3, 4), "int8"))
|
|
|
|
|
|
def test_collapse_sum_like_infer_ty_wrong_input_type():
|
|
bb = relax.BlockBuilder()
|
|
x0 = relax.Var("x", R.Tensor((3, 4, 5), "float32"))
|
|
x1 = relax.Var("x", relax.ShapeType((4, 5)))
|
|
x2 = relax.Var("x", relax.FuncType([], R.Tensor((2, 3, 4), "float32")))
|
|
|
|
with pytest.raises(TypeError):
|
|
bb.normalize(relax.op.collapse_sum_like(x0, x1))
|
|
|
|
with pytest.raises(TypeError):
|
|
bb.normalize(relax.op.collapse_sum_like(x2, x0))
|
|
|
|
|
|
def test_collapse_sum_like_infer_ty_shape_mismatch():
|
|
bb = relax.BlockBuilder()
|
|
x0 = relax.Var("x", R.Tensor((3, 4, 5), "float32"))
|
|
y0 = relax.Var("y", R.Tensor((3, 6, 5), "float32"))
|
|
a = tirx.Var("a", "int64")
|
|
b = tirx.Var("b", "int64")
|
|
x1 = relax.Var("z", R.Tensor((3, a, 5), "float32"))
|
|
y1 = relax.Var("w", R.Tensor((3, b, 5), "float32"))
|
|
|
|
s0 = relax.Var("s0", relax.ShapeType((3, 4, 5)))
|
|
s1 = relax.Var("s1", relax.ShapeType((3, 6, 5)))
|
|
x2 = relax.Var("x", relax.TensorType(s0, "float32"))
|
|
y2 = relax.Var("y", relax.TensorType(s1, "float32"))
|
|
|
|
s2 = relax.Var("s2", relax.ShapeType((3, a, 5)))
|
|
s3 = relax.Var("s3", relax.ShapeType((3, b, 5)))
|
|
x3 = relax.Var("x", relax.TensorType(s2, "float32"))
|
|
y3 = relax.Var("y", relax.TensorType(s3, "float32"))
|
|
|
|
with pytest.raises(ValueError):
|
|
bb.normalize(relax.op.collapse_sum_like(x0, y0))
|
|
|
|
with pytest.raises(ValueError):
|
|
bb.normalize(relax.op.collapse_sum_like(x1, y1))
|
|
|
|
with pytest.raises(ValueError):
|
|
bb.normalize(relax.op.collapse_sum_like(x2, y2))
|
|
|
|
with pytest.raises(ValueError):
|
|
bb.normalize(relax.op.collapse_sum_like(x3, y3))
|
|
|
|
|
|
def test_collapse_sum_to_infer_ty():
|
|
bb = relax.BlockBuilder()
|
|
x0 = relax.Var("x", R.Tensor((2, 3, 4), "float32"))
|
|
x1 = relax.Var("x", R.Tensor("float32", ndim=3))
|
|
x2 = relax.Var("x", R.Tensor("float32"))
|
|
x3 = relax.Var("x", R.Tensor((2, 3, 4)))
|
|
x4 = relax.Var("x", R.Tensor(ndim=3))
|
|
x5 = relax.Var("x", R.Tensor())
|
|
|
|
_check_inference(bb, relax.op.collapse_sum_to(x0, (3, 4)), relax.TensorType((3, 4), "float32"))
|
|
_check_inference(bb, relax.op.collapse_sum_to(x1, (3, 4)), relax.TensorType((3, 4), "float32"))
|
|
_check_inference(bb, relax.op.collapse_sum_to(x2, (3, 4)), relax.TensorType((3, 4), "float32"))
|
|
_check_inference(bb, relax.op.collapse_sum_to(x3, (3, 4)), relax.TensorType((3, 4), ""))
|
|
_check_inference(bb, relax.op.collapse_sum_to(x4, (3, 4)), relax.TensorType((3, 4), ""))
|
|
_check_inference(bb, relax.op.collapse_sum_to(x5, (3, 4)), relax.TensorType((3, 4), ""))
|
|
|
|
|
|
def test_collapse_sum_to_infer_ty_shape_symbolic():
|
|
bb = relax.BlockBuilder()
|
|
a = tirx.Var("a", "int64")
|
|
b = tirx.Var("b", "int64")
|
|
x0 = relax.Var("x", R.Tensor((3, 4, a), "float32"))
|
|
x1 = relax.Var("x", R.Tensor((3, 4, b + a), "float32"))
|
|
|
|
_check_inference(bb, relax.op.collapse_sum_to(x0, (4, a)), relax.TensorType((4, a), "float32"))
|
|
_check_inference(
|
|
bb, relax.op.collapse_sum_to(x1, (1, a + b)), relax.TensorType((1, a + b), "float32")
|
|
)
|
|
|
|
|
|
def test_collapse_sum_to_infer_ty_shape_var():
|
|
bb = relax.BlockBuilder()
|
|
s0 = relax.Var("s0", relax.ShapeType((2, 3, 4)))
|
|
s1 = relax.Var("s1", relax.ShapeType(ndim=3))
|
|
s2 = relax.Var("s2", relax.ShapeType())
|
|
x0 = relax.Var("x", relax.TensorType(s0, "float32"))
|
|
x1 = relax.Var("x", relax.TensorType(s1, "float32"))
|
|
x2 = relax.Var("x", relax.TensorType(s2, "float32"))
|
|
_check_inference(bb, relax.op.collapse_sum_to(x0, (3, 4)), relax.TensorType((3, 4), "float32"))
|
|
_check_inference(bb, relax.op.collapse_sum_to(x1, (3, 4)), relax.TensorType((3, 4), "float32"))
|
|
_check_inference(bb, relax.op.collapse_sum_to(x1, (3, 4)), relax.TensorType((3, 4), "float32"))
|
|
|
|
|
|
def test_collapse_sum_to_infer_ty_more_input_dtype():
|
|
bb = relax.BlockBuilder()
|
|
x0 = relax.Var("x", R.Tensor((2, 3, 4), "float16"))
|
|
x1 = relax.Var("x", R.Tensor((2, 3, 4), "int8"))
|
|
|
|
_check_inference(bb, relax.op.collapse_sum_to(x0, (3, 4)), relax.TensorType((3, 4), "float16"))
|
|
_check_inference(bb, relax.op.collapse_sum_to(x1, (3, 4)), relax.TensorType((3, 4), "int8"))
|
|
|
|
|
|
def test_collapse_sum_to_infer_ty_wrong_input_type():
|
|
bb = relax.BlockBuilder()
|
|
x0 = relax.Var("x", R.Tensor((3, 4, 5), "float32"))
|
|
x1 = relax.Var("x", relax.ShapeType((4, 5)))
|
|
x2 = relax.Var("x", relax.FuncType([], R.Tensor((2, 3, 4), "float32")))
|
|
|
|
with pytest.raises(TypeError):
|
|
bb.normalize(relax.op.collapse_sum_to(x0, x0))
|
|
|
|
with pytest.raises(TypeError):
|
|
bb.normalize(relax.op.collapse_sum_to(x0, x2))
|
|
|
|
with pytest.raises(TypeError):
|
|
bb.normalize(relax.op.collapse_sum_to(x1, x1))
|
|
|
|
|
|
def test_collapse_sum_to_infer_ty_shape_mismatch():
|
|
bb = relax.BlockBuilder()
|
|
x0 = relax.Var("x", R.Tensor((3, 4, 5), "float32"))
|
|
a = tirx.Var("a", "int64")
|
|
b = tirx.Var("b", "int64")
|
|
x1 = relax.Var("x", R.Tensor((3, a, 5), "float32"))
|
|
|
|
s0 = relax.Var("s0", relax.ShapeType((3, 4, 5)))
|
|
x2 = relax.Var("x", relax.TensorType(s0, "float32"))
|
|
|
|
s1 = relax.Var("s1", relax.ShapeType((3, a, 5)))
|
|
x3 = relax.Var("x", relax.TensorType(s1, "float32"))
|
|
|
|
with pytest.raises(ValueError):
|
|
bb.normalize(relax.op.collapse_sum_to(x0, (4, 4, 5)))
|
|
|
|
with pytest.raises(ValueError):
|
|
bb.normalize(relax.op.collapse_sum_to(x1, (3, b, 5)))
|
|
|
|
with pytest.raises(ValueError):
|
|
bb.normalize(relax.op.collapse_sum_to(x2, (4, 4, 5)))
|
|
|
|
with pytest.raises(ValueError):
|
|
bb.normalize(relax.op.collapse_sum_to(x3, (3, b, 5)))
|
|
|
|
|
|
def test_collapse_sum_to_infer_type_tgt_shape_var():
|
|
bb = relax.BlockBuilder()
|
|
a = tirx.Var("a", "int64")
|
|
b = tirx.Var("b", "int64")
|
|
c = tirx.Var("c", "int64")
|
|
d = tirx.Var("d", "int64")
|
|
s0 = relax.Var("s0", relax.ShapeType((3, a, b)))
|
|
s1 = relax.Var("s1", relax.ShapeType(ndim=3))
|
|
s2 = relax.Var("s2", relax.ShapeType())
|
|
x0 = relax.Var("x", R.Tensor((3, a, b), "float32"))
|
|
x1 = relax.Var("x", R.Tensor("float32", ndim=3))
|
|
x2 = relax.Var("x", R.Tensor(""))
|
|
x3 = relax.Var("x", relax.TensorType(s0, "float32"))
|
|
x4 = relax.Var("x", relax.TensorType(s1, "float32"))
|
|
x5 = relax.Var("x", relax.TensorType(s2, "float32"))
|
|
stgt0 = relax.Var("stgt0", relax.ShapeType((a, b)))
|
|
stgt1 = relax.Var("stgt1", relax.ShapeType(ndim=2))
|
|
stgt2 = relax.Var("stgt2", relax.ShapeType())
|
|
|
|
_check_inference(bb, relax.op.collapse_sum_to(x0, stgt0), relax.TensorType(stgt0, "float32"))
|
|
_check_inference(bb, relax.op.collapse_sum_to(x1, stgt0), relax.TensorType(stgt0, "float32"))
|
|
_check_inference(bb, relax.op.collapse_sum_to(x2, stgt0), relax.TensorType(stgt0, ""))
|
|
_check_inference(bb, relax.op.collapse_sum_to(x3, stgt0), relax.TensorType(stgt0, "float32"))
|
|
_check_inference(bb, relax.op.collapse_sum_to(x4, stgt0), relax.TensorType(stgt0, "float32"))
|
|
_check_inference(bb, relax.op.collapse_sum_to(x5, stgt0), relax.TensorType(stgt0, "float32"))
|
|
_check_inference(bb, relax.op.collapse_sum_to(x0, stgt1), relax.TensorType(stgt1, "float32"))
|
|
_check_inference(bb, relax.op.collapse_sum_to(x1, stgt1), relax.TensorType(stgt1, "float32"))
|
|
_check_inference(bb, relax.op.collapse_sum_to(x2, stgt1), relax.TensorType(stgt1, ""))
|
|
_check_inference(bb, relax.op.collapse_sum_to(x3, stgt1), relax.TensorType(stgt1, "float32"))
|
|
_check_inference(bb, relax.op.collapse_sum_to(x4, stgt1), relax.TensorType(stgt1, "float32"))
|
|
_check_inference(bb, relax.op.collapse_sum_to(x5, stgt1), relax.TensorType(stgt1, "float32"))
|
|
_check_inference(bb, relax.op.collapse_sum_to(x0, stgt2), relax.TensorType(stgt2, "float32"))
|
|
_check_inference(bb, relax.op.collapse_sum_to(x1, stgt2), relax.TensorType(stgt2, "float32"))
|
|
_check_inference(bb, relax.op.collapse_sum_to(x2, stgt2), relax.TensorType(stgt2, ""))
|
|
_check_inference(bb, relax.op.collapse_sum_to(x3, stgt2), relax.TensorType(stgt2, "float32"))
|
|
_check_inference(bb, relax.op.collapse_sum_to(x4, stgt2), relax.TensorType(stgt2, "float32"))
|
|
_check_inference(bb, relax.op.collapse_sum_to(x5, stgt2), relax.TensorType(stgt2, "float32"))
|
|
|
|
|
|
def test_repeat_infer_ty():
|
|
bb = relax.BlockBuilder()
|
|
vdev0 = VDevice("llvm")
|
|
x0 = relax.Var("x", R.Tensor((2, 10, 4), "float32"))
|
|
x1 = relax.Var("x", R.Tensor("float32", ndim=3))
|
|
x2 = relax.Var("x", R.Tensor("float32"))
|
|
x3 = relax.Var("x", R.Tensor((2, 10, 4)))
|
|
x4 = relax.Var("x", R.Tensor(ndim=3))
|
|
x5 = relax.Var("x", R.Tensor())
|
|
x6 = relax.Var("x", R.Tensor((2, 10, 4), "float32", vdev0))
|
|
|
|
_check_inference(
|
|
bb,
|
|
relax.op.repeat(x0, 2, axis=0),
|
|
relax.TensorType((4, 10, 4), "float32"),
|
|
)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.repeat(x6, 2, axis=0),
|
|
relax.TensorType((4, 10, 4), "float32", vdev0),
|
|
)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.repeat(x0, 2, axis=-2),
|
|
relax.TensorType((2, 20, 4), "float32"),
|
|
)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.repeat(x0, 2),
|
|
relax.TensorType((160,), "float32"),
|
|
)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.repeat(x1, 2, axis=0),
|
|
relax.TensorType(dtype="float32", ndim=3),
|
|
)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.repeat(x1, 2),
|
|
relax.TensorType(dtype="float32", ndim=1),
|
|
)
|
|
_check_inference(bb, relax.op.repeat(x2, 2, axis=0), relax.TensorType(dtype="float32"))
|
|
_check_inference(bb, relax.op.repeat(x2, 2), relax.TensorType(dtype="float32", ndim=1))
|
|
_check_inference(
|
|
bb,
|
|
relax.op.repeat(x3, 2, axis=0),
|
|
relax.TensorType((4, 10, 4), dtype=""),
|
|
)
|
|
_check_inference(bb, relax.op.repeat(x4, 2, axis=0), relax.TensorType(dtype="", ndim=3))
|
|
_check_inference(bb, relax.op.repeat(x5, 2, axis=0), relax.TensorType(dtype=""))
|
|
|
|
|
|
def test_repeat_infer_ty_shape_symbolic():
|
|
bb = relax.BlockBuilder()
|
|
a = tirx.Var("a", "int64")
|
|
b = tirx.Var("b", "int64")
|
|
c = tirx.Var("c", "int64")
|
|
x = relax.Var("x", R.Tensor((a, b, c), "float32"))
|
|
|
|
_check_inference(bb, relax.op.repeat(x, 2, 0), relax.TensorType((a * 2, b, c), "float32"))
|
|
_check_inference(
|
|
bb,
|
|
relax.op.repeat(x, 2, -1),
|
|
relax.TensorType((a, b, c * 2), "float32"),
|
|
)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.repeat(x, 2),
|
|
relax.TensorType((a * b * c * 2,), "float32"),
|
|
)
|
|
|
|
|
|
def test_repeat_infer_ty_more_input_dtype():
|
|
bb = relax.BlockBuilder()
|
|
x0 = relax.Var("x", R.Tensor((2, 3, 4), "float16"))
|
|
x1 = relax.Var("x", R.Tensor((2, 3, 4), "int8"))
|
|
|
|
_check_inference(bb, relax.op.repeat(x0, 2, 0), relax.TensorType((4, 3, 4), "float16"))
|
|
_check_inference(bb, relax.op.repeat(x1, 2, 0), relax.TensorType((4, 3, 4), "int8"))
|
|
|
|
|
|
def test_repeat_infer_ty_axis_out_of_range():
|
|
bb = relax.BlockBuilder()
|
|
x0 = relax.Var("x", R.Tensor((2, 10, 4), "float32"))
|
|
x1 = relax.Var("x", R.Tensor("float32", ndim=3))
|
|
x2 = relax.Var("x", R.Tensor("float32"))
|
|
|
|
with pytest.raises(ValueError):
|
|
bb.normalize(relax.op.repeat(x0, 2, 3))
|
|
with pytest.raises(ValueError):
|
|
bb.normalize(relax.op.repeat(x0, 2, -4))
|
|
with pytest.raises(ValueError):
|
|
bb.normalize(relax.op.repeat(x1, 2, 3))
|
|
with pytest.raises(ValueError):
|
|
bb.normalize(relax.op.repeat(x1, 2, -4))
|
|
# okay
|
|
bb.normalize(relax.op.repeat(x2, 2, 3))
|
|
bb.normalize(relax.op.repeat(x2, 2, -4))
|
|
|
|
|
|
def test_repeat_return_data_ty():
|
|
bb = relax.BlockBuilder()
|
|
x0 = relax.Var("x", R.Tensor((2, 10, 4), "float32"))
|
|
x1 = relax.Var("x", R.Tensor("float32", ndim=3))
|
|
x2 = relax.Var("x", R.Tensor("float32"))
|
|
|
|
_check_inference(bb, relax.op.repeat(x0, 1, 0), x0.ty)
|
|
_check_inference(bb, relax.op.repeat(x0, 1, -1), x0.ty)
|
|
_check_inference(bb, relax.op.repeat(x1, 1, 0), x1.ty)
|
|
_check_inference(bb, relax.op.repeat(x2, 1, 0), x2.ty)
|
|
|
|
|
|
def test_repeat_infer_ty_wrong_input_type():
|
|
bb = relax.BlockBuilder()
|
|
x0 = relax.Var("x", relax.ShapeType((2, 3, 4, 5)))
|
|
x1 = relax.Var("x", relax.FuncType([], R.Tensor((2, 3, 4, 5), "float32")))
|
|
x2 = relax.Var("x", R.Tensor((2, 3, 4, 5), "float32"))
|
|
r1 = tirx.Var("r", "float32")
|
|
r2 = tirx.StringImm("abc")
|
|
|
|
with pytest.raises(TypeError):
|
|
bb.normalize(relax.op.repeat(x0, 2))
|
|
with pytest.raises(TypeError):
|
|
bb.normalize(relax.op.repeat(x1, 2))
|
|
with pytest.raises(TypeError):
|
|
bb.normalize(relax.op.repeat(x2, 1.5))
|
|
with pytest.raises(TypeError):
|
|
bb.normalize(relax.op.repeat(x2, r1))
|
|
with pytest.raises(TypeError):
|
|
bb.normalize(relax.op.repeat(x2, r2))
|
|
|
|
|
|
def test_tile_infer_ty():
|
|
bb = relax.BlockBuilder()
|
|
vdev0 = VDevice("llvm")
|
|
x0 = relax.Var("x", R.Tensor((2, 10, 4), "float32"))
|
|
x1 = relax.Var("x", R.Tensor("float32", ndim=3))
|
|
x2 = relax.Var("x", R.Tensor("float32"))
|
|
x3 = relax.Var("x", R.Tensor((2, 10, 4)))
|
|
x4 = relax.Var("x", R.Tensor(ndim=3))
|
|
x5 = relax.Var("x", R.Tensor())
|
|
x6 = relax.Var("x", R.Tensor((2, 10, 4), "float32", vdev0))
|
|
|
|
_check_inference(
|
|
bb,
|
|
relax.op.tile(x0, 2),
|
|
relax.TensorType((2, 10, 8), "float32"),
|
|
)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.tile(x6, 2),
|
|
relax.TensorType((2, 10, 8), "float32", vdev0),
|
|
)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.tile(x0, (3, 2)),
|
|
relax.TensorType((2, 30, 8), "float32"),
|
|
)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.tile(x0, (4, 3, 2)),
|
|
relax.TensorType((8, 30, 8), "float32"),
|
|
)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.tile(x0, (5, 4, 3, 2)),
|
|
relax.TensorType((5, 8, 30, 8), "float32"),
|
|
)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.tile(x1, 2),
|
|
relax.TensorType(dtype="float32", ndim=3),
|
|
)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.tile(x1, (5, 4, 3, 2)),
|
|
relax.TensorType(dtype="float32", ndim=4),
|
|
)
|
|
_check_inference(bb, relax.op.tile(x2, (5, 4, 3, 2)), relax.TensorType(dtype="float32"))
|
|
_check_inference(
|
|
bb,
|
|
relax.op.tile(x3, 2),
|
|
relax.TensorType((2, 10, 8), dtype=""),
|
|
)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.tile(x3, (5, 4, 3, 2)),
|
|
relax.TensorType((5, 8, 30, 8), dtype=""),
|
|
)
|
|
_check_inference(bb, relax.op.tile(x4, 2), relax.TensorType(dtype="", ndim=3))
|
|
_check_inference(bb, relax.op.tile(x4, (5, 4, 3, 2)), relax.TensorType(dtype="", ndim=4))
|
|
_check_inference(bb, relax.op.tile(x5, (5, 4, 3, 2)), relax.TensorType(dtype=""))
|
|
|
|
|
|
def test_tile_infer_ty_shape_symbolic():
|
|
bb = relax.BlockBuilder()
|
|
a = tirx.Var("a", "int64")
|
|
b = tirx.Var("b", "int64")
|
|
c = tirx.Var("c", "int64")
|
|
x = relax.Var("x", R.Tensor((a, b, c), "float32"))
|
|
|
|
_check_inference(bb, relax.op.tile(x, 2), relax.TensorType((a, b, c * 2), "float32"))
|
|
_check_inference(bb, relax.op.tile(x, (3, 2)), relax.TensorType((a, b * 3, c * 2), "float32"))
|
|
_check_inference(
|
|
bb, relax.op.tile(x, (4, 3, 2)), relax.TensorType((a * 4, b * 3, c * 2), "float32")
|
|
)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.tile(x, (5, 4, 3, 2)),
|
|
relax.TensorType((5, a * 4, b * 3, c * 2), "float32"),
|
|
)
|
|
|
|
|
|
def test_tile_infer_ty_more_input_dtype():
|
|
bb = relax.BlockBuilder()
|
|
x0 = relax.Var("x", R.Tensor((2, 3, 4), "float16"))
|
|
x1 = relax.Var("x", R.Tensor((2, 3, 4), "int8"))
|
|
|
|
_check_inference(bb, relax.op.tile(x0, (3, 2)), relax.TensorType((2, 9, 8), "float16"))
|
|
_check_inference(bb, relax.op.tile(x1, (3, 2)), relax.TensorType((2, 9, 8), "int8"))
|
|
|
|
|
|
def test_tile_return_data_ty():
|
|
bb = relax.BlockBuilder()
|
|
x0 = relax.Var("x", R.Tensor((2, 10, 4), "float32"))
|
|
x1 = relax.Var("x", R.Tensor("float32", ndim=3))
|
|
x2 = relax.Var("x", R.Tensor("float32"))
|
|
|
|
_check_inference(bb, relax.op.tile(x0, 1), x0.ty)
|
|
_check_inference(bb, relax.op.tile(x0, (1, 1)), x0.ty)
|
|
_check_inference(bb, relax.op.tile(x0, (1, 1, 1)), x0.ty)
|
|
_check_inference(bb, relax.op.tile(x1, 1), x1.ty)
|
|
_check_inference(bb, relax.op.tile(x2, 1), x2.ty)
|
|
|
|
|
|
def test_tile_infer_ty_wrong_input_type():
|
|
bb = relax.BlockBuilder()
|
|
x0 = relax.Var("x", relax.ShapeType((2, 3, 4, 5)))
|
|
x1 = relax.Var("x", relax.FuncType([], R.Tensor((2, 3, 4, 5), "float32")))
|
|
x2 = relax.Var("x", R.Tensor((2, 3, 4, 5), "float32"))
|
|
r1 = tirx.Var("a", "float32")
|
|
r2 = tirx.StringImm("abc")
|
|
|
|
with pytest.raises(TypeError):
|
|
bb.normalize(relax.op.tile(x0, 2))
|
|
with pytest.raises(TypeError):
|
|
bb.normalize(relax.op.tile(x1, 2))
|
|
with pytest.raises(TypeError):
|
|
bb.normalize(relax.op.tile(x2, (2, 1.5, 2)))
|
|
with pytest.raises(TypeError):
|
|
bb.normalize(relax.op.tile(x2, (2, r1)))
|
|
with pytest.raises(TypeError):
|
|
bb.normalize(relax.op.tile(x2, r2))
|
|
|
|
|
|
def test_flip_infer_ty():
|
|
bb = relax.BlockBuilder()
|
|
vdev0 = VDevice("llvm")
|
|
x0 = relax.Var("x", R.Tensor((2, 10, 4), "float32"))
|
|
x1 = relax.Var("x", R.Tensor("float16", ndim=3))
|
|
x2 = relax.Var("x", R.Tensor("int32"))
|
|
x3 = relax.Var("x", R.Tensor((2, 10, 4)))
|
|
x4 = relax.Var("x", R.Tensor(ndim=3))
|
|
x5 = relax.Var("x", R.Tensor((2, 10, 4), "float32", vdev0))
|
|
|
|
_check_inference(bb, relax.op.flip(x0, axis=1), relax.TensorType((2, 10, 4), "float32"))
|
|
_check_inference(bb, relax.op.flip(x5, axis=1), relax.TensorType((2, 10, 4), "float32", vdev0))
|
|
_check_inference(bb, relax.op.flip(x1, axis=0), R.Tensor("float16", ndim=3))
|
|
_check_inference(bb, relax.op.flip(x2, axis=0), R.Tensor("int32"))
|
|
_check_inference(bb, relax.op.flip(x3, axis=2), R.Tensor((2, 10, 4)))
|
|
_check_inference(bb, relax.op.flip(x4, axis=2), R.Tensor(ndim=3))
|
|
|
|
|
|
def test_flip_infer_ty_shape_symbolic():
|
|
bb = relax.BlockBuilder()
|
|
a = tirx.Var("a", "int64")
|
|
b = tirx.Var("b", "int64")
|
|
x = relax.Var("x", R.Tensor((a, b), "float32"))
|
|
|
|
_check_inference(bb, relax.op.flip(x, axis=0), relax.TensorType((a, b), "float32"))
|
|
|
|
|
|
def test_flip_infer_ty_wrong_inputs():
|
|
bb = relax.BlockBuilder()
|
|
x0 = relax.Var("x", R.Tensor((2, 10, 4), "float32"))
|
|
|
|
with pytest.raises(ValueError):
|
|
bb.normalize(relax.op.flip(x0, axis=3))
|
|
|
|
|
|
def test_reverse_sequence_infer_ty():
|
|
bb = relax.BlockBuilder()
|
|
vdev0 = VDevice("llvm")
|
|
x0 = relax.Var("x", R.Tensor((2, 10, 4), "float32"))
|
|
x1 = relax.Var("x", R.Tensor("float16", ndim=3))
|
|
x2 = relax.Var("x", R.Tensor("int32"))
|
|
x3 = relax.Var("x", R.Tensor((2, 10, 4)))
|
|
x4 = relax.Var("x", R.Tensor(ndim=3))
|
|
x5 = relax.Var("x", R.Tensor((2, 10, 4), "float32", vdev0))
|
|
s0 = relax.Var("s", R.Tensor((2,), "int32"))
|
|
s1 = relax.Var("s", R.Tensor("int64", ndim=1))
|
|
|
|
_check_inference(
|
|
bb,
|
|
relax.op.reverse_sequence(x0, s0, seq_axis=1, batch_axis=0),
|
|
relax.TensorType((2, 10, 4), "float32"),
|
|
)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.reverse_sequence(x5, s0, seq_axis=1, batch_axis=0),
|
|
relax.TensorType((2, 10, 4), "float32", vdev0),
|
|
)
|
|
_check_inference(
|
|
bb, relax.op.reverse_sequence(x1, s1, seq_axis=0, batch_axis=1), R.Tensor("float16", ndim=3)
|
|
)
|
|
_check_inference(bb, relax.op.reverse_sequence(x2, s1), R.Tensor("int32"))
|
|
_check_inference(bb, relax.op.reverse_sequence(x3, s0), R.Tensor((2, 10, 4)))
|
|
_check_inference(bb, relax.op.reverse_sequence(x4, s1), R.Tensor(ndim=3))
|
|
|
|
|
|
def test_reverse_sequence_infer_ty_shape_symbolic():
|
|
bb = relax.BlockBuilder()
|
|
a = tirx.Var("a", "int64")
|
|
b = tirx.Var("b", "int64")
|
|
x = relax.Var("x", R.Tensor((a, b, 4), "float32"))
|
|
seq_lengths = relax.Var("seq_lengths", R.Tensor((b,), "int64"))
|
|
|
|
_check_inference(
|
|
bb,
|
|
relax.op.reverse_sequence(x, seq_lengths, seq_axis=0, batch_axis=1),
|
|
relax.TensorType((a, b, 4), "float32"),
|
|
)
|
|
|
|
|
|
def test_reverse_sequence_infer_ty_wrong_inputs():
|
|
bb = relax.BlockBuilder()
|
|
x = relax.Var("x", R.Tensor((2, 10, 4), "float32"))
|
|
seq_lengths = relax.Var("seq_lengths", R.Tensor((2,), "int32"))
|
|
seq_lengths_2d = relax.Var("seq_lengths", R.Tensor((2, 1), "int32"))
|
|
seq_lengths_float = relax.Var("seq_lengths", R.Tensor((2,), "float32"))
|
|
seq_lengths_int16 = relax.Var("seq_lengths", R.Tensor((2,), "int16"))
|
|
seq_lengths_mismatch = relax.Var("seq_lengths", R.Tensor((3,), "int32"))
|
|
not_tensor = relax.Var("seq_lengths", relax.AnyType())
|
|
|
|
with pytest.raises(TypeError):
|
|
bb.normalize(relax.op.reverse_sequence(x, not_tensor))
|
|
with pytest.raises(ValueError):
|
|
bb.normalize(relax.op.reverse_sequence(x, seq_lengths_2d))
|
|
with pytest.raises(ValueError):
|
|
bb.normalize(relax.op.reverse_sequence(x, seq_lengths_float))
|
|
with pytest.raises(ValueError):
|
|
bb.normalize(relax.op.reverse_sequence(x, seq_lengths_int16))
|
|
with pytest.raises(ValueError):
|
|
bb.normalize(relax.op.reverse_sequence(x, seq_lengths, seq_axis=3))
|
|
with pytest.raises(ValueError):
|
|
bb.normalize(relax.op.reverse_sequence(x, seq_lengths, batch_axis=-4))
|
|
with pytest.raises(ValueError):
|
|
bb.normalize(relax.op.reverse_sequence(x, seq_lengths_mismatch))
|
|
|
|
|
|
def test_gather_elements_infer_ty():
|
|
bb = relax.BlockBuilder()
|
|
vdev0 = VDevice("llvm")
|
|
x0 = relax.Var("x", R.Tensor((2, 3, 4), "float32"))
|
|
x1 = relax.Var("x", R.Tensor("float32", ndim=3))
|
|
x2 = relax.Var("x", R.Tensor("float32"))
|
|
x3 = relax.Var("x", R.Tensor((2, 3, 4), "float32", vdev0))
|
|
i0 = relax.Var("i", R.Tensor((2, 3, 4), "int64"))
|
|
i1 = relax.Var("i", R.Tensor((2, 3, 4)))
|
|
i2 = relax.Var("i", R.Tensor("int64", ndim=3))
|
|
i3 = relax.Var("i", R.Tensor(ndim=3))
|
|
i4 = relax.Var("i", R.Tensor((2, 3, 4), "int64", vdev0))
|
|
|
|
_check_inference(
|
|
bb, relax.op.gather_elements(x0, i0, axis=1), relax.TensorType((2, 3, 4), "float32")
|
|
)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.gather_elements(x3, i4, axis=1),
|
|
relax.TensorType((2, 3, 4), "float32", vdev0),
|
|
)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.gather_elements(x1, i0, axis=1),
|
|
relax.TensorType((2, 3, 4), dtype="float32"),
|
|
)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.gather_elements(x2, i0, axis=0),
|
|
relax.TensorType(dtype="float32", ndim=-1),
|
|
)
|
|
_check_inference(
|
|
bb, relax.op.gather_elements(x0, i1, axis=1), relax.TensorType((2, 3, 4), "float32")
|
|
)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.gather_elements(x1, i2, axis=1),
|
|
relax.TensorType(dtype="float32", ndim=3),
|
|
)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.gather_elements(x2, i3, axis=0),
|
|
relax.TensorType(dtype="float32", ndim=-1),
|
|
)
|
|
|
|
|
|
def test_gather_elements_infer_ty_shape_symbolic():
|
|
bb = relax.BlockBuilder()
|
|
a = tirx.Var("a", "int64")
|
|
b = tirx.Var("b", "int64")
|
|
x = relax.Var("x", R.Tensor((a, b), "float32"))
|
|
i = relax.Var("i", R.Tensor((a, b), "int64"))
|
|
|
|
_check_inference(
|
|
bb, relax.op.gather_elements(x, i, axis=1), relax.TensorType((a, b), "float32")
|
|
)
|
|
|
|
|
|
def test_gather_elements_infer_ty_wrong_inputs():
|
|
bb = relax.BlockBuilder()
|
|
x0 = relax.Var("x", R.Tensor((2, 3, 4), "float32"))
|
|
x1 = relax.Var("x", R.Tensor((2, 3), "float32"))
|
|
i0 = relax.Var("i", R.Tensor((2, 3, 4), "int64"))
|
|
i1 = relax.Var("i", R.Tensor((2, 3), "int64"))
|
|
i2 = relax.Var("i", R.Tensor((2, 3, 4), "float32"))
|
|
|
|
with pytest.raises(ValueError):
|
|
bb.normalize(relax.op.gather_elements(x0, i0, axis=3))
|
|
with pytest.raises(ValueError):
|
|
bb.normalize(relax.op.gather_elements(x0, i1))
|
|
with pytest.raises(ValueError):
|
|
bb.normalize(relax.op.gather_elements(x1, i0))
|
|
with pytest.raises(TypeError):
|
|
bb.normalize(relax.op.gather_elements(x0, i2))
|
|
|
|
|
|
def test_gather_nd_infer_ty():
|
|
bb = relax.BlockBuilder()
|
|
vdev0 = VDevice("llvm")
|
|
x0 = relax.Var("x", R.Tensor((2, 3, 4), "float32"))
|
|
x1 = relax.Var("x", R.Tensor("float32", ndim=3))
|
|
x2 = relax.Var("x", R.Tensor("float32"))
|
|
x3 = relax.Var("x", R.Tensor((2, 3, 4), "float32", vdev0))
|
|
i0 = relax.Var("i", R.Tensor((2, 2), "int64"))
|
|
i1 = relax.Var("i", R.Tensor((2, 2)))
|
|
i2 = relax.Var("i", R.Tensor("int64", ndim=2))
|
|
i3 = relax.Var("i", R.Tensor(ndim=2))
|
|
i4 = relax.Var("i", R.Tensor((2, 2), "int64", vdev0))
|
|
|
|
_check_inference(bb, relax.op.gather_nd(x0, i0), relax.TensorType((2, 4), "float32"))
|
|
_check_inference(bb, relax.op.gather_nd(x3, i4), relax.TensorType((2, 4), "float32", vdev0))
|
|
_check_inference(bb, relax.op.gather_nd(x1, i0), relax.TensorType(dtype="float32", ndim=2))
|
|
_check_inference(bb, relax.op.gather_nd(x2, i0), relax.TensorType(dtype="float32", ndim=-1))
|
|
_check_inference(bb, relax.op.gather_nd(x0, i1), relax.TensorType((2, 4), "float32"))
|
|
_check_inference(bb, relax.op.gather_nd(x1, i2), relax.TensorType(dtype="float32"))
|
|
_check_inference(bb, relax.op.gather_nd(x2, i3), relax.TensorType(dtype="float32"))
|
|
|
|
|
|
def test_gather_nd_infer_ty_shape_symbolic():
|
|
bb = relax.BlockBuilder()
|
|
a = tirx.Var("a", "int64")
|
|
b = tirx.Var("b", "int64")
|
|
c = tirx.Var("c", "int64")
|
|
x = relax.Var("x", R.Tensor((a, b, c), "float32"))
|
|
i = relax.Var("i", R.Tensor((2, 2), "int64"))
|
|
|
|
_check_inference(bb, relax.op.gather_nd(x, i), relax.TensorType((2, c), "float32"))
|
|
|
|
|
|
def test_gather_nd_infer_ty_wrong_inputs():
|
|
bb = relax.BlockBuilder()
|
|
x0 = relax.Var("x", R.Tensor((2, 3, 4), "float32"))
|
|
i0 = relax.Var("i", R.Tensor((2, 4), "int64")) # indices too long
|
|
i1 = relax.Var("i", R.Tensor((2, 2), "float32")) # wrong dtype
|
|
|
|
with pytest.raises(ValueError):
|
|
bb.normalize(relax.op.gather_nd(x0, i0))
|
|
with pytest.raises(TypeError):
|
|
bb.normalize(relax.op.gather_nd(x0, i1))
|
|
|
|
|
|
def test_scatter_elements_infer_ty():
|
|
bb = relax.BlockBuilder()
|
|
vdev0 = VDevice("llvm")
|
|
d0 = relax.Var("data", R.Tensor((4, 4), "float32"))
|
|
d1 = relax.Var("data", R.Tensor(dtype="float32", ndim=2))
|
|
d2 = relax.Var("data", R.Tensor("float32"))
|
|
d3 = relax.Var("data", R.Tensor((4, 4), "float32", vdev0))
|
|
i0 = relax.Var("indices", R.Tensor((2, 2), "int64"))
|
|
i1 = relax.Var("indices", R.Tensor((2, 2)))
|
|
i2 = relax.Var("indices", R.Tensor(dtype="int64", ndim=2))
|
|
i3 = relax.Var("indices", R.Tensor(ndim=2))
|
|
i4 = relax.Var("indices", R.Tensor((2, 2), "int64", vdev0))
|
|
u0 = relax.Var("updates", R.Tensor((2, 2), "float32"))
|
|
u1 = relax.Var("updates", R.Tensor((2, 2), "float32", vdev0))
|
|
_check_inference(
|
|
bb,
|
|
relax.op.scatter_elements(d0, i0, u0, 0, "updates"),
|
|
relax.TensorType((4, 4), dtype="float32"),
|
|
)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.scatter_elements(d3, i4, u1, 0, "updates"),
|
|
relax.TensorType((4, 4), dtype="float32", vdevice=vdev0),
|
|
)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.scatter_elements(d1, i0, u0, 0, "updates"),
|
|
relax.TensorType(dtype="float32", ndim=2),
|
|
)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.scatter_elements(d2, i0, u0, 0, "updates"),
|
|
relax.TensorType(dtype="float32", ndim=-1),
|
|
)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.scatter_elements(d0, i1, u0, 0, "updates"),
|
|
relax.TensorType((4, 4), dtype="float32"),
|
|
)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.scatter_elements(d1, i1, u0, 0, "updates"),
|
|
relax.TensorType(dtype="float32", ndim=2),
|
|
)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.scatter_elements(d2, i1, u0, 0, "updates"),
|
|
relax.TensorType(dtype="float32", ndim=-1),
|
|
)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.scatter_elements(d0, i2, u0, 0, "updates"),
|
|
relax.TensorType((4, 4), dtype="float32"),
|
|
)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.scatter_elements(d1, i2, u0, 0, "updates"),
|
|
relax.TensorType(dtype="float32", ndim=2),
|
|
)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.scatter_elements(d2, i2, u0, 0, "updates"),
|
|
relax.TensorType(dtype="float32", ndim=-1),
|
|
)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.scatter_elements(d0, i3, u0, 0, "updates"),
|
|
relax.TensorType((4, 4), dtype="float32"),
|
|
)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.scatter_elements(d1, i3, u0, 0, "updates"),
|
|
relax.TensorType(dtype="float32", ndim=2),
|
|
)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.scatter_elements(d2, i3, u0, 0, "updates"),
|
|
relax.TensorType(dtype="float32", ndim=-1),
|
|
)
|
|
# Test with unknown dtype for data
|
|
d_unknown = relax.Var("data", R.Tensor((4, 4)))
|
|
_check_inference(
|
|
bb,
|
|
relax.op.scatter_elements(d_unknown, i0, u0, 0, "updates"),
|
|
relax.TensorType((4, 4), dtype=""),
|
|
)
|
|
# Test with unknown dtype for updates
|
|
u_unknown = relax.Var("updates", R.Tensor((2, 2)))
|
|
_check_inference(
|
|
bb,
|
|
relax.op.scatter_elements(d0, i0, u_unknown, 0, "updates"),
|
|
relax.TensorType((4, 4), dtype="float32"),
|
|
)
|
|
|
|
|
|
def test_scatter_elements_infer_ty_symbolic_shape():
|
|
bb = relax.BlockBuilder()
|
|
a = tirx.Var("a", "int64")
|
|
b = tirx.Var("b", "int64")
|
|
c = tirx.Var("c", "int64")
|
|
d = tirx.Var("d", "int64")
|
|
e = tirx.Var("e", "int64")
|
|
f = tirx.Var("f", "int64")
|
|
|
|
d0 = relax.Var("data", R.Tensor((a, b), "float32"))
|
|
i0 = relax.Var("indices", R.Tensor((c, d), "int64"))
|
|
u0 = relax.Var("updates", R.Tensor((c, d), "float32"))
|
|
u1 = relax.Var("updates", R.Tensor((e, f), "float32"))
|
|
|
|
_check_inference(
|
|
bb,
|
|
relax.op.scatter_elements(d0, i0, u0, 0, "updates"),
|
|
relax.TensorType((a, b), dtype="float32"),
|
|
)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.scatter_elements(d0, i0, u1, 0, "updates"),
|
|
relax.TensorType((a, b), dtype="float32"),
|
|
)
|
|
|
|
|
|
def test_scatter_elements_infer_ty_wrong_indices_type():
|
|
bb = relax.BlockBuilder()
|
|
d0 = relax.Var("data", R.Tensor((4, 4), "float32"))
|
|
i0 = relax.Var("indices", R.Tensor((2, 2), "float32"))
|
|
u0 = relax.Var("updates", R.Tensor((2, 2), "float32"))
|
|
|
|
with pytest.raises(TypeError):
|
|
bb.normalize(relax.op.scatter_elements(d0, i0, u0))
|
|
|
|
|
|
def test_scatter_elements_infer_ty_rank_shape_mismatch():
|
|
a = tirx.Var("a", "int64")
|
|
b = tirx.Var("b", "int64")
|
|
|
|
bb = relax.BlockBuilder()
|
|
d0 = relax.Var("data", R.Tensor((4, 4), "float32"))
|
|
i0 = relax.Var("indices", R.Tensor((3, 3), "int64"))
|
|
i1 = relax.Var("indices", R.Tensor((3, 3, 3), "int64"))
|
|
i2 = relax.Var("indices", R.Tensor((a, b), "int64"))
|
|
u0 = relax.Var("updates", R.Tensor((3, 2), "float32"))
|
|
u1 = relax.Var("updates", R.Tensor((3, 2, 3), "float32"))
|
|
u2 = relax.Var("updates", R.Tensor((3, 3, 3), "float32"))
|
|
u3 = relax.Var("updates", R.Tensor((a + 1, b), "float32"))
|
|
u4 = relax.Var("updates", R.Tensor((3, 3), "float16"))
|
|
|
|
with pytest.raises(ValueError):
|
|
bb.normalize(relax.op.scatter_elements(d0, i0, u0))
|
|
with pytest.raises(ValueError):
|
|
bb.normalize(relax.op.scatter_elements(d0, i1, u0))
|
|
with pytest.raises(ValueError):
|
|
bb.normalize(relax.op.scatter_elements(d0, i0, u1))
|
|
with pytest.raises(ValueError):
|
|
bb.normalize(relax.op.scatter_elements(d0, i1, u1))
|
|
with pytest.raises(ValueError):
|
|
bb.normalize(relax.op.scatter_elements(d0, i1, u2))
|
|
with pytest.raises(ValueError):
|
|
bb.normalize(relax.op.scatter_elements(d0, i2, u3))
|
|
with pytest.raises(TypeError):
|
|
bb.normalize(relax.op.scatter_elements(d0, i0, u4))
|
|
|
|
|
|
def test_scatter_nd_infer_ty():
|
|
bb = relax.BlockBuilder()
|
|
|
|
d0 = relax.Var("data", R.Tensor((8,), "float32"))
|
|
i0 = relax.Var("indices", R.Tensor((4, 1), "int64"))
|
|
u0 = relax.Var("updates", R.Tensor((4,), "float32"))
|
|
|
|
_check_inference(
|
|
bb,
|
|
relax.op.scatter_nd(d0, i0, u0, "update"),
|
|
relax.TensorType((8,), dtype="float32"),
|
|
)
|
|
|
|
d1 = relax.Var("data", R.Tensor((4, 4, 4), "float32"))
|
|
i1 = relax.Var("indices", R.Tensor((2, 1), "int64"))
|
|
u1 = relax.Var("updates", R.Tensor((2, 4, 4), "float32"))
|
|
|
|
_check_inference(
|
|
bb,
|
|
relax.op.scatter_nd(d1, i1, u1, "update"),
|
|
relax.TensorType((4, 4, 4), dtype="float32"),
|
|
)
|
|
|
|
|
|
def test_meshgrid_infer_ty():
|
|
bb = relax.BlockBuilder()
|
|
vdev0 = VDevice("llvm")
|
|
t0 = relax.Var("t0", R.Tensor((3,), "float32"))
|
|
t1 = relax.Var("t1", R.Tensor((4,), "float32"))
|
|
t2 = relax.Var("t2", R.Tensor("float32", ndim=1))
|
|
t3 = relax.Var("t3", R.Tensor((5,), "float32", vdev0))
|
|
|
|
_check_inference(
|
|
bb,
|
|
relax.op.meshgrid((t0, t1), indexing="ij"),
|
|
relax.TupleType([relax.TensorType((3, 4), "float32"), relax.TensorType((3, 4), "float32")]),
|
|
)
|
|
|
|
_check_inference(
|
|
bb,
|
|
relax.op.meshgrid((t3, t1), indexing="ij"),
|
|
relax.TupleType(
|
|
[
|
|
relax.TensorType((5, 4), "float32", vdev0),
|
|
relax.TensorType((5, 4), "float32", vdev0),
|
|
]
|
|
),
|
|
)
|
|
|
|
_check_inference(
|
|
bb,
|
|
relax.op.meshgrid((t2, t1), indexing="xy"),
|
|
relax.TupleType(
|
|
[
|
|
relax.TensorType(dtype="float32", ndim=2),
|
|
relax.TensorType(dtype="float32", ndim=2),
|
|
]
|
|
),
|
|
)
|
|
|
|
_check_inference(
|
|
bb,
|
|
relax.op.meshgrid((t0,), indexing="ij"),
|
|
relax.TupleType([relax.TensorType((3,), "float32")]),
|
|
)
|
|
|
|
|
|
def test_one_hot_infer_ty():
|
|
bb = relax.BlockBuilder()
|
|
|
|
# Test case 1: Basic usage
|
|
i0 = relax.Var("indices", R.Tensor((3,), "int32"))
|
|
_check_inference(
|
|
bb,
|
|
relax.op.one_hot(i0, tirx.FloatImm("float32", 1.0), tirx.FloatImm("float32", 0.0), 5),
|
|
relax.TensorType((3, 5), "float32"),
|
|
)
|
|
|
|
# Test case 2: With specified axis
|
|
i1 = relax.Var("indices", R.Tensor((2, 2), "int32"))
|
|
_check_inference(
|
|
bb,
|
|
relax.op.one_hot(i1, tirx.IntImm("int64", 1), tirx.IntImm("int64", 0), 3, axis=1),
|
|
relax.TensorType((2, 3, 2), "int64"),
|
|
)
|
|
|
|
# Test case 3: With symbolic shape
|
|
n = tirx.Var("n", "int64")
|
|
i2 = relax.Var("indices", R.Tensor((n,), "int32"))
|
|
_check_inference(
|
|
bb,
|
|
relax.op.one_hot(i2, tirx.FloatImm("float32", 1.0), tirx.FloatImm("float32", 0.0), 4),
|
|
relax.TensorType((n, 4), "float32"),
|
|
)
|
|
|
|
# Test case 4: With unknown shape
|
|
i3 = relax.Var("indices", R.Tensor("int32"))
|
|
_check_inference(
|
|
bb,
|
|
relax.op.one_hot(i3, tirx.FloatImm("float32", 1.0), tirx.FloatImm("float32", 0.0), 6),
|
|
relax.TensorType(dtype="float32"),
|
|
)
|
|
|
|
# Test case 5: With different on_value and off_value dtypes
|
|
i3 = relax.Var("indices", R.Tensor((2, 3), "int32"))
|
|
with pytest.raises(tvm.error.InternalError):
|
|
bb.normalize(
|
|
relax.op.one_hot(i3, tirx.FloatImm("float32", 1.0), tirx.IntImm("int64", 0), 5)
|
|
)
|
|
|
|
# Test case 6: With invalid indices dtype
|
|
i4 = relax.Var("indices", R.Tensor((2, 3), "float32"))
|
|
with pytest.raises(TypeError):
|
|
bb.normalize(
|
|
relax.op.one_hot(i4, tirx.FloatImm("float32", 1.0), tirx.FloatImm("float32", 0.0), 5)
|
|
)
|
|
|
|
# Test case 7: With invalid depth
|
|
i5 = relax.Var("indices", R.Tensor((2, 3), "int32"))
|
|
with pytest.raises(tvm.error.InternalError):
|
|
bb.normalize(
|
|
relax.op.one_hot(i5, tirx.FloatImm("float32", 1.0), tirx.FloatImm("float32", 0.0), -1)
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
tvm.testing.main()
|