chore: import upstream snapshot with attribution

tests/python/relax/test_frontend_nn_tensor.py

@@ -0,0 +1,165 @@
+# 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: E501, F401
+import numpy as np
+import pytest
+
+import tvm
+import tvm.testing
+from tvm import relax
+from tvm.relax.frontend.nn import Module, Tensor, spec
+from tvm.script import relax as R
+
+
+def test_tensor_from_numpy():
+    x = np.random.rand(1, 10)
+    tensor_x = Tensor.from_const(x)
+    assert tensor_x.shape == [1, 10]
+    assert tensor_x.ndim == 2
+    assert tensor_x.dtype == "float32"
+    assert repr(tensor_x) == 'Tensor([1, 10], "float32")'
+
+
+def test_tensor_from_scalar():
+    x = 123.321
+    tensor_x = Tensor.from_scalar(x, dtype="float16")
+    assert tensor_x.shape == []
+    assert tensor_x.ndim == 0
+    assert tensor_x.dtype == "float16"
+    assert repr(tensor_x) == 'Tensor([], "float16")'
+
+
+def test_tensor_op_binary_tensor_tensor():
+    class Model(Module):
+        def test(self, x: Tensor, y: Tensor):
+            z0 = x + y
+            z1 = x * y
+            z2 = x / y
+            z3 = x.maximum(y)
+            z4 = x.minimum(y)
+            return (z0, z1, z2, z3, z4)
+
+    # fmt: off
+    @R.function
+    def test(x: R.Tensor((1, 10), dtype="float32"), y: R.Tensor((2, 1), dtype="float32"), _io: R.Any) -> R.Tuple(R.Tuple(R.Tensor((2, 10), dtype="float32"), R.Tensor((2, 10), dtype="float32"), R.Tensor((2, 10), dtype="float32"), R.Tensor((2, 10), dtype="float32"), R.Tensor((2, 10), dtype="float32")), R.Tuple(R.Any)):
+        R.func_attr({"num_input": 3})
+        with R.dataflow():
+            add: R.Tensor((2, 10), dtype="float32") = R.add(x, y)
+            mul: R.Tensor((2, 10), dtype="float32") = R.multiply(x, y)
+            divide: R.Tensor((2, 10), dtype="float32") = R.divide(x, y)
+            maximum: R.Tensor((2, 10), dtype="float32") = R.maximum(x, y)
+            minimum: R.Tensor((2, 10), dtype="float32") = R.minimum(x, y)
+            gv1: R.Tuple(R.Tuple(R.Tensor((2, 10), dtype="float32"), R.Tensor((2, 10), dtype="float32"), R.Tensor((2, 10), dtype="float32"), R.Tensor((2, 10), dtype="float32"), R.Tensor((2, 10), dtype="float32")), R.Tuple(R.Any)) = (add, mul, divide, maximum, minimum), (_io,)
+            R.output(gv1)
+        return gv1
+    # fmt: on
+
+    m = Model()
+    irmodule, _ = m.export_tvm(
+        spec={"test": {"x": spec.Tensor([1, 10], "float32"), "y": spec.Tensor([2, 1], "float32")}},
+        debug=True,
+    )
+
+    tvm.ir.assert_structural_equal(irmodule["test"], test)
+
+
+def test_tensor_op_binary_tensor_scalar():
+    class Model(Module):
+        def test(self, x: Tensor):
+            y = 10
+            z0 = x + y
+            z1 = y + x
+            z2 = x * y
+            z3 = x / y
+            z4 = x.maximum(y)
+            z5 = x.minimum(y)
+            return (z0, z1, z2, z3, z4, z5)
+
+    # fmt: off
+    @R.function
+    def test(x: R.Tensor((1, 10), dtype="float32"), _io: R.Any) -> R.Tuple(R.Tuple(R.Tensor((1, 10), dtype="float32"), R.Tensor((1, 10), dtype="float32"), R.Tensor((1, 10), dtype="float32"), R.Tensor((1, 10), dtype="float32"), R.Tensor((1, 10), dtype="float32"), R.Tensor((1, 10), dtype="float32")), R.Tuple(R.Any)):
+        R.func_attr({"num_input": 2})
+        with R.dataflow():
+            add: R.Tensor((1, 10), dtype="float32") = R.add(x, R.const(10, "float32"))
+            add1: R.Tensor((1, 10), dtype="float32") = R.add(x, R.const(10, "float32"))
+            mul: R.Tensor((1, 10), dtype="float32") = R.multiply(x, R.const(10, "float32"))
+            divide: R.Tensor((1, 10), dtype="float32") = R.divide(x, R.const(10, "float32"))
+            maximum: R.Tensor((1, 10), dtype="float32") = R.maximum(x, R.const(10, "float32"))
+            minimum: R.Tensor((1, 10), dtype="float32") = R.minimum(x, R.const(10, "float32"))
+            gv1: R.Tuple(R.Tuple(R.Tensor((1, 10), dtype="float32"), R.Tensor((1, 10), dtype="float32"), R.Tensor((1, 10), dtype="float32"), R.Tensor((1, 10), dtype="float32"), R.Tensor((1, 10), dtype="float32"), R.Tensor((1, 10), dtype="float32")), R.Tuple(R.Any)) = (add, add1, mul, divide, maximum, minimum), (_io,)
+            R.output(gv1)
+        return gv1
+    # fmt: on
+
+    m = Model()
+    irmodule, _ = m.export_tvm(spec={"test": {"x": spec.Tensor([1, 10], "float32")}}, debug=True)
+
+    tvm.ir.assert_structural_equal(irmodule["test"], test)
+
+
+def test_tensor_op_datatype():
+    class Model(Module):
+        def test(self, x: Tensor):
+            z0 = x.astype(dtype="float16")
+            return z0
+
+    # fmt: off
+    @R.function
+    def test(x: R.Tensor((1, 10), dtype="float32"), _io: R.Any) -> R.Tuple(R.Tensor((1, 10), dtype="float16"), R.Tuple(R.Any)):
+        R.func_attr({"num_input": 2})
+        with R.dataflow():
+            astype: R.Tensor((1, 10), dtype="float16") = R.astype(x, dtype="float16")
+            gv1: R.Tuple(R.Tensor((1, 10), dtype="float16"), R.Tuple(R.Any)) = astype, (_io,)
+            R.output(gv1)
+        return gv1
+    # fmt: on
+
+    m = Model()
+    irmodule, _ = m.export_tvm(spec={"test": {"x": spec.Tensor([1, 10], "float32")}}, debug=True)
+
+    tvm.ir.assert_structural_equal(irmodule["test"], test)
+
+
+def test_tensor_op_manipulate():
+    class Model(Module):
+        def test(self, x: Tensor):
+            z0 = x.reshape(2, 5, 2)
+            z1 = x.permute_dims(2, 1, 0)
+            z2 = x.repeat(2, axis=1)
+            return (z0, z1, z2)
+
+    # fmt: off
+    @R.function
+    def test(x: R.Tensor((2, 1, 10), dtype="float32"), _io: R.Any) -> R.Tuple(R.Tuple(R.Tensor((2, 5, 2), dtype="float32"), R.Tensor((10, 1, 2), dtype="float32"), R.Tensor((2, 2, 10), dtype="float32")), R.Tuple(R.Any)):
+        R.func_attr({"num_input": 2})
+        with R.dataflow():
+            reshape: R.Tensor((2, 5, 2), dtype="float32") = R.reshape(x, R.shape([2, 5, 2]))
+            permute_dims: R.Tensor((10, 1, 2), dtype="float32") = R.permute_dims(x, axes=[2, 1, 0])
+            repeat: R.Tensor((2, 2, 10), dtype="float32") = R.repeat(x, repeats=2, axis=1)
+            gv1: R.Tuple(R.Tuple(R.Tensor((2, 5, 2), dtype="float32"), R.Tensor((10, 1, 2), dtype="float32"), R.Tensor((2, 2, 10), dtype="float32")), R.Tuple(R.Any)) = (reshape, permute_dims, repeat), (_io,)
+            R.output(gv1)
+        return gv1
+    # fmt: on
+
+    m = Model()
+    irmodule, _ = m.export_tvm(spec={"test": {"x": spec.Tensor([2, 1, 10], "float32")}}, debug=True)
+
+    tvm.ir.assert_structural_equal(irmodule["test"], test)
+
+
+if __name__ == "__main__":
+    tvm.testing.main()