chore: import upstream snapshot with attribution

tests/python/ir/test_node_reflection.py

@@ -0,0 +1,186 @@
+# 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: E712, F401, F841
+import json
+import sys
+
+import numpy as np
+import pytest
+import tvm_ffi
+
+import tvm
+import tvm.testing
+from tvm import te
+
+
+def test_const_saveload_json():
+    # save load json
+    x = tvm.tirx.const(1, "int32")
+    y = tvm.tirx.const(10, "int32")
+    z = x + y
+    z = z + z
+    json_str = tvm.ir.save_json(z)
+    zz = tvm.ir.load_json(json_str)
+    tvm.ir.assert_structural_equal(zz, z, map_free_vars=True)
+
+
+def test_save_json_metadata_version():
+    obj = tvm.runtime.convert([1, 2])
+    json_str = tvm.ir.save_json(obj)
+    assert json.loads(json_str)["metadata"]["tvm_version"] == tvm.__version__
+    assert list(tvm.ir.load_json(json_str)) == [1, 2]
+
+
+def _test_infinity_value(value, dtype):
+    x = tvm.tirx.const(value, dtype)
+    json_str = tvm.ir.save_json(x)
+    tvm.ir.assert_structural_equal(x, tvm.ir.load_json(json_str))
+
+
+def test_infinity_value():
+    _test_infinity_value(float("inf"), "float64")
+    _test_infinity_value(float("-inf"), "float64")
+    _test_infinity_value(float("inf"), "float32")
+    _test_infinity_value(float("-inf"), "float32")
+
+
+def _test_minmax_value(value):
+    json_str = tvm.ir.save_json(value)
+    tvm.ir.assert_structural_equal(value, tvm.ir.load_json(json_str))
+
+
+def test_minmax_value():
+    _test_minmax_value(tvm.tirx.min_value("float32"))
+    _test_minmax_value(tvm.tirx.max_value("float32"))
+
+
+def test_make_smap():
+    # save load json
+    x = tvm.tirx.const(1, "int32")
+    y = tvm.tirx.const(10, "int32")
+    z = tvm.tirx.Add(x, y)
+    smap = tvm.runtime.convert({"z": z, "x": x})
+    json_str = tvm.ir.save_json(tvm.runtime.convert([smap]))
+    arr = tvm.ir.load_json(json_str)
+    assert len(arr) == 1
+    assert arr[0]["z"].a == arr[0]["x"]
+    tvm.ir.assert_structural_equal(arr, [smap], map_free_vars=True)
+
+
+def test_make_node():
+    x = tvm.ir.make_node("ir.IntImm", ty=tvm.ir.PrimType("int32"), value=10, span=None)
+    assert isinstance(x, tvm.tirx.IntImm)
+    assert x.value == 10
+    A = te.placeholder((10,), name="A")
+    AA = tvm.ir.make_node(
+        "te.Tensor", shape=A.shape, dtype=A.dtype, op=A.op, value_index=A.value_index
+    )
+    assert AA.op == A.op
+    assert AA.value_index == A.value_index
+
+    y = tvm.ir.make_node(
+        "ir.IntImm", ty=tvm.ir.PrimType(tvm_ffi.core.String("int32")), value=10, span=None
+    )
+    assert isinstance(y, tvm.tirx.IntImm)
+    assert y.value == 10
+
+
+def test_make_sum():
+    A = te.placeholder((2, 10), name="A")
+    k = te.reduce_axis((0, 10), "k")
+    B = te.compute((2,), lambda i: te.sum(A[i, k], axis=k), name="B")
+    json_str = tvm.ir.save_json(B)
+    BB = tvm.ir.load_json(json_str)
+    assert B.op.body[0].combiner is not None
+    assert BB.op.body[0].combiner is not None
+
+
+def test_string():
+    # non printable str, need to store by b64
+    s1 = tvm_ffi.core.String("xy\x01z")
+    s2 = tvm.ir.load_json(tvm.ir.save_json(s1))
+    tvm.ir.assert_structural_equal(s1, s2)
+
+    # printable str, need to store by repr_str
+    s1 = tvm_ffi.core.String("xyz")
+    s2 = tvm.ir.load_json(tvm.ir.save_json(s1))
+    tvm.ir.assert_structural_equal(s1, s2)
+
+
+def test_pass_config():
+    cfg = tvm.transform.PassContext(
+        opt_level=1,
+        config={
+            "tirx.UnrollLoop": {
+                "auto_max_step": 10,
+            }
+        },
+    )
+    cfg.opt_level == 1
+
+    assert cfg.config["tirx.UnrollLoop"].auto_max_step == 10
+    # default option
+    assert cfg.config["tirx.UnrollLoop"].explicit_unroll == True
+
+    # schema checking for specific config key
+    with pytest.raises(TypeError):
+        cfg = tvm.transform.PassContext(config={"tirx.UnrollLoop": {"invalid": 1}})
+
+    # schema check for un-registered config
+    with pytest.raises(AttributeError):
+        cfg = tvm.transform.PassContext(config={"inavlid-opt": True})
+
+    # schema check for wrong type
+    with pytest.raises(AttributeError):
+        cfg = tvm.transform.PassContext(config={"tirx.UnrollLoop": 1})
+
+
+def test_dict():
+    x = tvm.tirx.const(1)  # a class that has Python-defined methods
+    # instances should see the full class dict
+    assert set(dir(x.__class__)) <= set(dir(x))
+
+
+def test_tensor():
+    dev = tvm.cpu(0)
+    tvm_arr = tvm.runtime.tensor(np.random.rand(4), device=dev)
+    tvm_arr2 = tvm.ir.load_json(tvm.ir.save_json(tvm_arr))
+    tvm.ir.assert_structural_equal(tvm_arr, tvm_arr2)
+    np.testing.assert_array_equal(tvm_arr.numpy(), tvm_arr2.numpy())
+
+
+def test_tensor_dict():
+    dev = tvm.cpu(0)
+    m1 = {
+        "key1": tvm.runtime.tensor(np.random.rand(4), device=dev),
+        "key2": tvm.runtime.tensor(np.random.rand(4), device=dev),
+    }
+    m2 = tvm.ir.load_json(tvm.ir.save_json(m1))
+    tvm.ir.assert_structural_equal(m1, m2)
+
+
+def test_free_var_equal():
+    x = tvm.tirx.Var("x", dtype="int32")
+    y = tvm.tirx.Var("y", dtype="int32")
+    z = tvm.tirx.Var("z", dtype="int32")
+    v1 = x + y
+    v1 = y + z
+    tvm.ir.assert_structural_equal(x, z, map_free_vars=True)
+
+
+if __name__ == "__main__":
+    tvm.testing.main()