chore: import upstream snapshot with attribution

tests/python/relax/test_bind_params.py

@@ -0,0 +1,150 @@
+# 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.
+
+
+import numpy as np
+import pytest
+
+import tvm
+import tvm.script
+import tvm.testing
+from tvm import relax, tirx
+from tvm.script import relax as R
+
+param_specification = tvm.testing.parameter("by_string", "by_var")
+param_shape = tvm.testing.parameter("static_shape", "dynamic_shape", "ndim", "arbitrary")
+tensor_param_dtype = tvm.testing.parameter("float32", None)
+
+
+def test_bind_tensor_param(param_specification, param_shape, tensor_param_dtype):
+    if param_shape == "static_shape":
+        shape = [16]
+        ndim = -1
+    elif param_shape == "dynamic_shape":
+        shape = [tirx.Var("N", "int64")]
+        ndim = -1
+    elif param_shape == "ndim":
+        shape = None
+        ndim = 1
+    elif param_shape == "arbitrary":
+        shape = None
+        ndim = -1
+    else:
+        raise ValueError(f"Unknown param_shape: {param_shape}")
+
+    @R.function
+    def before(A: R.Tensor(shape, ndim=ndim, dtype=tensor_param_dtype)):
+        R.func_attr({"global_symbol": "main"})
+        B: R.Tensor(shape=shape, ndim=ndim, dtype=tensor_param_dtype) = A
+        out = R.add(B, B)
+        return out
+
+    np_data = np.arange(16).astype("float32")
+    inlined_relax_const = relax.const(np_data)
+
+    @R.function
+    def expected() -> R.Tensor([16], "float32"):
+        R.func_attr({"global_symbol": "main"})
+        B = inlined_relax_const
+        out = R.add(B, B)
+        return out
+
+    if param_specification == "by_string":
+        var = "A"
+    elif param_specification == "by_var":
+        var = before.params[0]
+    else:
+        raise ValueError("Unknown param_specification: {param_specification}")
+
+    after = before.bind_params({var: np.arange(16).astype("float32")})
+
+    tvm.ir.assert_structural_equal(expected, after)
+
+
+def test_bind_shape_param(param_shape):
+    if param_shape == "static_shape":
+        shape = [16]
+        ndim = -1
+    elif param_shape == "dynamic_shape":
+        shape = [tirx.Var("N", "int64")]
+        ndim = -1
+    elif param_shape == "ndim":
+        shape = None
+        ndim = 1
+    elif param_shape == "arbitrary":
+        shape = None
+        ndim = -1
+    else:
+        raise ValueError(f"Unknown param_shape: {param_shape}")
+
+    @R.function
+    def before(A: R.Shape(shape, ndim=ndim)):
+        R.func_attr({"global_symbol": "main"})
+        B: R.Shape(shape, ndim=ndim) = A
+        return B
+
+    @R.function
+    def expected() -> R.Shape([16]):
+        R.func_attr({"global_symbol": "main"})
+        B = R.ShapeExpr([16])
+        return B
+
+    after = before.bind_params({"A": relax.ShapeExpr([16])})
+
+    tvm.ir.assert_structural_equal(expected, after)
+
+
+prim_value_dtype = tvm.testing.parameter("int64", "int32", "float32")
+
+
+def test_bind_prim_value(prim_value_dtype):
+    prim_type = tvm.ir.PrimType(prim_value_dtype)
+    param = relax.Var("A", prim_type)
+    before = relax.Function([param], param, prim_type).with_attr("global_symbol", "main")
+    value = tirx.const(16, prim_value_dtype)
+
+    after = before.bind_params({"A": value})
+
+    assert not after.params
+    tvm.ir.assert_structural_equal(after.ret_ty, prim_type)
+    tvm.ir.assert_structural_equal(after.body.body, value)
+
+
+def test_error_on_unknown_var():
+    @R.function
+    def before(A: R.Tensor([16], dtype="float32")):
+        R.func_attr({"global_symbol": "main"})
+        return A
+
+    unknown_var = relax.Var("unknown_var")
+
+    with pytest.raises(RuntimeError):
+        before.bind_params({unknown_var: np.arange(16).astype("float32")})
+
+
+def test_error_on_unknown_var_name():
+    @R.function
+    def before(A: R.Tensor([16], dtype="float32")):
+        R.func_attr({"global_symbol": "main"})
+        return A
+
+    with pytest.raises(RuntimeError):
+        before.bind_params({"unknown_var_name": np.arange(16).astype("float32")})
+
+
+if __name__ == "__main__":
+    tvm.testing.main()