# 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()