This is an automated email from the ASF dual-hosted git repository.
tqchen pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm-ffi.git
The following commit(s) were added to refs/heads/main by this push:
new 6adc8df7 [EXTRA] Introduce StructuralKey (#453)
6adc8df7 is described below
commit 6adc8df7d2180ea14e463d3beeaa3d7eecb6f897
Author: Tianqi Chen <[email protected]>
AuthorDate: Mon Feb 16 21:08:20 2026 -0500
[EXTRA] Introduce StructuralKey (#453)
This PR introduces StructuralKey as a convenient wrapper class to
indicate use of structural equal/hash in comparisons. It can be useful
to serve as an intermediate wrapper class in structural equality related
checks.
---
include/tvm/ffi/extra/structural_key.h | 103 +++++++++++++++
python/tvm_ffi/__init__.py | 12 ++
python/tvm_ffi/_ffi_api.py | 6 +
python/tvm_ffi/structural.py | 230 +++++++++++++++++++++++++++++++++
src/ffi/extra/reflection_extra.cc | 38 +++++-
src/ffi/extra/structural_equal.cc | 4 +-
tests/cpp/extra/test_structural_key.cc | 79 +++++++++++
tests/python/test_structural.py | 92 +++++++++++++
8 files changed, 562 insertions(+), 2 deletions(-)
diff --git a/include/tvm/ffi/extra/structural_key.h
b/include/tvm/ffi/extra/structural_key.h
new file mode 100644
index 00000000..08afd018
--- /dev/null
+++ b/include/tvm/ffi/extra/structural_key.h
@@ -0,0 +1,103 @@
+/*
+ * 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.
+ */
+/*!
+ * \file tvm/ffi/extra/structural_key.h
+ * \brief Structural-key wrapper that caches structural hash.
+ */
+#ifndef TVM_FFI_EXTRA_STRUCTURAL_KEY_H_
+#define TVM_FFI_EXTRA_STRUCTURAL_KEY_H_
+
+#include <tvm/ffi/any.h>
+#include <tvm/ffi/extra/structural_equal.h>
+#include <tvm/ffi/extra/structural_hash.h>
+
+#include <cstddef>
+#include <cstdint>
+#include <functional>
+#include <utility>
+
+namespace tvm {
+namespace ffi {
+
+/*!
+ * \brief Object node that contains a key and its cached structural hash.
+ */
+class StructuralKeyObj : public Object {
+ public:
+ /*! \brief The key value. */
+ Any key;
+ /*! \brief Cached structural hash of \p key, encoded as int64_t. */
+ int64_t hash_i64{0};
+
+ // Default constructor to support reflection-based initialization.
+ StructuralKeyObj() = default;
+ /*!
+ * \brief Construct a StructuralKeyObj from a key and cache its structural
hash.
+ * \param key The key value.
+ */
+ explicit StructuralKeyObj(Any key)
+ : key(std::move(key)),
hash_i64(static_cast<int64_t>(StructuralHash::Hash(this->key))) {}
+
+ /// \cond Doxygen_Suppress
+ TVM_FFI_DECLARE_OBJECT_INFO_FINAL("ffi.StructuralKey", StructuralKeyObj,
Object);
+ /// \endcond
+};
+
+/*!
+ * \brief ObjectRef wrapper for StructuralKeyObj.
+ *
+ * StructuralKey can be used to ensure that we use structural equality and
hash for the wrapped key.
+ */
+class StructuralKey : public ObjectRef {
+ public:
+ /*!
+ * \brief Construct a StructuralKey from a key.
+ * \param key The key value.
+ */
+ explicit StructuralKey(Any key) :
ObjectRef(make_object<StructuralKeyObj>(std::move(key))) {}
+
+ bool operator==(const StructuralKey& other) const {
+ if (this->same_as(other)) {
+ return true;
+ }
+ if (this->get()->hash_i64 != other->hash_i64) {
+ return false;
+ }
+ return StructuralEqual::Equal(this->get()->key, other->key);
+ }
+ bool operator!=(const StructuralKey& other) const { return !(*this ==
other); }
+
+ /// \cond Doxygen_Suppress
+ TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(StructuralKey, ObjectRef,
StructuralKeyObj);
+ /// \endcond
+};
+
+} // namespace ffi
+} // namespace tvm
+
+namespace std {
+template <>
+struct hash<tvm::ffi::StructuralKey> {
+ size_t operator()(const tvm::ffi::StructuralKey& key) const {
+ return static_cast<size_t>(static_cast<uint64_t>(key->hash_i64));
+ }
+};
+} // namespace std
+
+#endif // TVM_FFI_EXTRA_STRUCTURAL_KEY_H_
diff --git a/python/tvm_ffi/__init__.py b/python/tvm_ffi/__init__.py
index 21eca3b6..3eedbc55 100644
--- a/python/tvm_ffi/__init__.py
+++ b/python/tvm_ffi/__init__.py
@@ -54,9 +54,16 @@ from ._tensor import from_dlpack, Tensor, Shape
from .container import Array, List, Map
from .module import Module, system_lib, load_module
from .stream import StreamContext, get_raw_stream, use_raw_stream,
use_torch_stream
+from .structural import (
+ StructuralKey,
+ get_first_structural_mismatch,
+ structural_equal,
+ structural_hash,
+)
from . import serialization
from . import access_path
from . import dataclasses
+from . import structural
from . import cpp
# optional module to speedup dlpack conversion
@@ -104,6 +111,7 @@ __all__ = [
"ObjectConvertible",
"Shape",
"StreamContext",
+ "StructuralKey",
"Tensor",
"__version__",
"__version_tuple__",
@@ -114,6 +122,7 @@ __all__ = [
"device",
"dtype",
"from_dlpack",
+ "get_first_structural_mismatch",
"get_global_func",
"get_global_func_metadata",
"get_raw_stream",
@@ -124,6 +133,9 @@ __all__ = [
"register_object",
"remove_global_func",
"serialization",
+ "structural",
+ "structural_equal",
+ "structural_hash",
"system_lib",
"use_raw_stream",
"use_torch_stream",
diff --git a/python/tvm_ffi/_ffi_api.py b/python/tvm_ffi/_ffi_api.py
index 92606c9e..bf77b76a 100644
--- a/python/tvm_ffi/_ffi_api.py
+++ b/python/tvm_ffi/_ffi_api.py
@@ -84,6 +84,9 @@ if TYPE_CHECKING:
def ModuleWriteToFile(_0: Module, _1: str, _2: str, /) -> None: ...
def Shape(*args: Any) -> Any: ...
def String(_0: str, /) -> str: ...
+ def StructuralKey(_0: Any, /) -> Any: ...
+ def StructuralKeyEqual(_0: Any, _1: Any, /) -> bool: ...
+ def StructuralEqual(_0: Any, _1: Any, _2: bool, _3: bool, /) -> bool: ...
def StructuralHash(_0: Any, _1: bool, _2: bool, /) -> int: ...
def SystemLib(*args: Any) -> Any: ...
def ToJSONGraph(_0: Any, _1: Any, /) -> Any: ...
@@ -143,7 +146,10 @@ __all__ = [
"ModuleWriteToFile",
"Shape",
"String",
+ "StructuralEqual",
"StructuralHash",
+ "StructuralKey",
+ "StructuralKeyEqual",
"SystemLib",
"ToJSONGraph",
"ToJSONGraphString",
diff --git a/python/tvm_ffi/structural.py b/python/tvm_ffi/structural.py
new file mode 100644
index 00000000..120268a1
--- /dev/null
+++ b/python/tvm_ffi/structural.py
@@ -0,0 +1,230 @@
+# 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.
+# pylint: disable=invalid-name
+"""Structural helper objects and functions."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+if TYPE_CHECKING:
+ from .access_path import AccessPath
+
+from . import _ffi_api
+from .core import Object
+from .registry import register_object
+
+__all__ = [
+ "StructuralKey",
+ "get_first_structural_mismatch",
+ "structural_equal",
+ "structural_hash",
+]
+
+
+def structural_equal(
+ lhs: Any, rhs: Any, map_free_vars: bool = False, skip_tensor_content: bool
= False
+) -> bool:
+ """Check structural equality between two values.
+
+ Structural equality compares the *shape/content structure* of two values
+ instead of Python object identity. For container-like values, this means
+ recursive comparison of elements/fields. For object types that provide
+ structural equal hooks, those hooks are used.
+
+ Parameters
+ ----------
+ lhs
+ Left-hand side value.
+ rhs
+ Right-hand side value.
+
+ map_free_vars
+ Whether free variables (variables without a definition site) can be
+ mapped to each other during comparison.
+
+ skip_tensor_content
+ Whether to skip tensor data content when comparing tensors.
+
+ Returns
+ -------
+ result
+ Whether the two values are structurally equal.
+
+ Examples
+ --------
+ .. code-block:: python
+
+ import tvm_ffi
+
+ assert tvm_ffi.structural_equal([1, 2, 3], [1, 2, 3])
+ assert not tvm_ffi.structural_equal([1, 2, 3], [1, 2, 4])
+
+ See Also
+ --------
+ :py:func:`tvm_ffi.structural_hash`
+ Hash function compatible with structural equality.
+ :py:func:`tvm_ffi.get_first_structural_mismatch`
+ Mismatch diagnostics with access paths.
+
+ """
+ return _ffi_api.StructuralEqual(lhs, rhs, map_free_vars,
skip_tensor_content)
+
+
+def structural_hash(
+ value: Any, map_free_vars: bool = False, skip_tensor_content: bool = False
+) -> int:
+ """Compute structural hash of a value.
+
+ This hash is designed to be consistent with :py:func:`structural_equal`
+ under the same options. If two values are structurally equal, they should
+ have the same structural hash.
+
+ Parameters
+ ----------
+ value
+ Input value to hash.
+
+ map_free_vars
+ Whether free variables mapped to each other during hashing.
+
+ skip_tensor_content
+ Whether tensor data content is ignored when hashing tensors.
+
+ Returns
+ -------
+ hash_value
+ Structural hash value.
+
+ Examples
+ --------
+ .. code-block:: python
+
+ import tvm_ffi
+
+ h0 = tvm_ffi.structural_hash([1, 2, 3])
+ h1 = tvm_ffi.structural_hash([1, 2, 3])
+ assert h0 == h1
+
+ Notes
+ -----
+ Structural hash is intended for hash-table bucketing and fast pre-checks.
+ Always use structural equality to confirm semantic equivalence.
+
+ """
+ # need to mask the result so negative values are converted to u64
+ # this is because hash value were stored as int64_t in the C++ code
+ return _ffi_api.StructuralHash(value, map_free_vars, skip_tensor_content)
& 0xFFFFFFFFFFFFFFFF
+
+
+def get_first_structural_mismatch(
+ lhs: Any, rhs: Any, map_free_vars: bool = False, skip_tensor_content: bool
= False
+) -> tuple[AccessPath, AccessPath] | None:
+ """Like structural_equal(), but returns the AccessPath pair of the first
detected mismatch.
+
+ Parameters
+ ----------
+ lhs
+ The left operand.
+
+ rhs
+ The right operand.
+
+ map_free_vars
+ Whether free variables (i.e. variables without a definition site)
should be mapped
+ as equal to each other.
+
+ skip_tensor_content
+ Whether to skip the data content of tensor.
+
+ Returns
+ -------
+ mismatch: tuple[AccessPath, AccessPath] | None
+ `None` if `lhs` and `rhs` are structurally equal.
+ Otherwise, a tuple of two AccessPath objects that point to the first
detected mismatch.
+
+ """
+ return _ffi_api.GetFirstStructuralMismatch(lhs, rhs, map_free_vars,
skip_tensor_content)
+
+
+@register_object("ffi.StructuralKey")
+class StructuralKey(Object):
+ """Hash-cached structural key wrapper.
+
+ This wrapper can be used to hint that a dict uses structural equality and
hash for the key.
+
+ Examples
+ --------
+ Use ``StructuralKey`` with Python dictionaries when you want key lookup by
+ structural semantics:
+
+ .. code-block:: python
+
+ import tvm_ffi
+
+ k0 = tvm_ffi.StructuralKey([1, 2, 3])
+ k1 = tvm_ffi.StructuralKey([1, 2, 3])
+ k2 = tvm_ffi.StructuralKey([1, 2, 4])
+
+ d = {k0: "value-a", k2: "value-b"}
+ assert d[k1] == "value-a" # k1 matches k0 structurally
+ assert d[k2] == "value-b"
+
+ It can also be used directly with :py:class:`tvm_ffi.Map`:
+
+ .. code-block:: python
+
+ m = tvm_ffi.Map({k0: 1, k1: 2})
+ assert len(m) == 1
+ assert m[k0] == 2
+
+ See Also
+ --------
+ :py:func:`tvm_ffi.structural_equal`
+ Structural equality comparison.
+ :py:func:`tvm_ffi.structural_hash`
+ Structural hash computation.
+
+ """
+
+ # tvm-ffi-stubgen(begin): object/ffi.StructuralKey
+ # fmt: off
+ key: Any
+ hash_i64: int
+ # fmt: on
+ # tvm-ffi-stubgen(end)
+
+ def __init__(self, key: Any) -> None:
+ """Create a structural key from ``key``.
+
+ Parameters
+ ----------
+ key
+ The underlying value used for structural hash/equality.
+
+ """
+ self.__init_handle_by_constructor__(_ffi_api.StructuralKey, key)
+
+ def __hash__(self) -> int:
+ """Return cached structural hash."""
+ # need to mask the result so negative values are converted to u64
+ # this is because hash value were stored as int64_t in the C++ code
+ return self.hash_i64 & 0xFFFFFFFFFFFFFFFF
+
+ def __eq__(self, other: Any) -> bool:
+ """Compare by structural equality."""
+ return isinstance(other, StructuralKey) and
_ffi_api.StructuralKeyEqual(self, other)
diff --git a/src/ffi/extra/reflection_extra.cc
b/src/ffi/extra/reflection_extra.cc
index 66994163..20190092 100644
--- a/src/ffi/extra/reflection_extra.cc
+++ b/src/ffi/extra/reflection_extra.cc
@@ -21,6 +21,7 @@
*
* \brief Extra reflection registrations. *
*/
+#include <tvm/ffi/extra/structural_key.h>
#include <tvm/ffi/reflection/access_path.h>
#include <tvm/ffi/reflection/accessor.h>
#include <tvm/ffi/reflection/registry.h>
@@ -133,11 +134,46 @@ inline void AccessPathRegisterReflection() {
[](const AccessPath& self, const AccessPath& other) { return
self->PathEqual(other); });
}
+int64_t StructuralKeyHash(const Any& key) {
+ const auto* key_obj = key.as<StructuralKeyObj>();
+ TVM_FFI_ICHECK_NOTNULL(key_obj);
+ return key_obj->hash_i64;
+}
+
+bool StructuralKeyEqual(const Any& lhs, const Any& rhs) {
+ if (lhs.same_as(rhs)) {
+ return true;
+ }
+ const auto* lhs_obj = lhs.as<StructuralKeyObj>();
+ const auto* rhs_obj = rhs.as<StructuralKeyObj>();
+ TVM_FFI_ICHECK_NOTNULL(lhs_obj);
+ TVM_FFI_ICHECK_NOTNULL(rhs_obj);
+ if (lhs_obj->hash_i64 != rhs_obj->hash_i64) {
+ return false;
+ }
+ return StructuralEqual::Equal(lhs_obj->key, rhs_obj->key);
+}
+
+inline void StructuralKeyRegisterReflection() {
+ namespace refl = tvm::ffi::reflection;
+ refl::ObjectDef<StructuralKeyObj>()
+ .def_ro("key", &StructuralKeyObj::key)
+ .def_ro("hash_i64", &StructuralKeyObj::hash_i64);
+ refl::TypeAttrDef<StructuralKeyObj>()
+ .attr("__any_hash__", reinterpret_cast<void*>(&StructuralKeyHash))
+ .attr("__any_equal__", reinterpret_cast<void*>(&StructuralKeyEqual));
+}
+
TVM_FFI_STATIC_INIT_BLOCK() {
namespace refl = tvm::ffi::reflection;
AccessStepRegisterReflection();
AccessPathRegisterReflection();
- refl::GlobalDef().def_packed("ffi.MakeObjectFromPackedArgs",
MakeObjectFromPackedArgs);
+ StructuralKeyRegisterReflection();
+
+ refl::GlobalDef()
+ .def_packed("ffi.MakeObjectFromPackedArgs", MakeObjectFromPackedArgs)
+ .def("ffi.StructuralKey", [](Any key) { return
StructuralKey(std::move(key)); })
+ .def("ffi.StructuralKeyEqual", StructuralKeyEqual);
}
} // namespace reflection
diff --git a/src/ffi/extra/structural_equal.cc
b/src/ffi/extra/structural_equal.cc
index 2fa05f24..3e8cb54a 100644
--- a/src/ffi/extra/structural_equal.cc
+++ b/src/ffi/extra/structural_equal.cc
@@ -454,7 +454,9 @@ Optional<reflection::AccessPathPair>
StructuralEqual::GetFirstMismatch(const Any
TVM_FFI_STATIC_INIT_BLOCK() {
namespace refl = tvm::ffi::reflection;
- refl::GlobalDef().def("ffi.GetFirstStructuralMismatch",
StructuralEqual::GetFirstMismatch);
+ refl::GlobalDef()
+ .def("ffi.StructuralEqual", StructuralEqual::Equal)
+ .def("ffi.GetFirstStructuralMismatch",
StructuralEqual::GetFirstMismatch);
// ensure the type attribute column is presented in the system even if it is
empty.
refl::EnsureTypeAttrColumn("__s_equal__");
}
diff --git a/tests/cpp/extra/test_structural_key.cc
b/tests/cpp/extra/test_structural_key.cc
new file mode 100644
index 00000000..bc6b646c
--- /dev/null
+++ b/tests/cpp/extra/test_structural_key.cc
@@ -0,0 +1,79 @@
+/*
+ * 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.
+ */
+#include <gtest/gtest.h>
+#include <tvm/ffi/container/array.h>
+#include <tvm/ffi/container/map.h>
+#include <tvm/ffi/extra/structural_hash.h>
+#include <tvm/ffi/extra/structural_key.h>
+
+#include <unordered_map>
+
+namespace {
+
+using namespace tvm::ffi;
+
+TEST(StructuralKey, EqualityAndStdHash) {
+ StructuralKey k1(Array<int>{1, 2, 3});
+ StructuralKey k2(Array<int>{1, 2, 3});
+ StructuralKey k3(Array<int>{1, 2, 4});
+
+ EXPECT_FALSE(k1.same_as(k2));
+ EXPECT_EQ(k1->hash_i64, static_cast<int64_t>(StructuralHash::Hash(k1->key)));
+ EXPECT_EQ(k2->hash_i64, static_cast<int64_t>(StructuralHash::Hash(k2->key)));
+
+ EXPECT_TRUE(k1 == k2);
+ EXPECT_FALSE(k1 != k2);
+ EXPECT_FALSE(k1 == k3);
+ EXPECT_TRUE(k1 != k3);
+ EXPECT_EQ(std::hash<StructuralKey>()(k1), std::hash<StructuralKey>()(k2));
+ EXPECT_NE(std::hash<StructuralKey>()(k1), std::hash<StructuralKey>()(k3));
+
+ std::unordered_map<StructuralKey, int> map;
+ map.emplace(k1, 10);
+ map[k2] = 20;
+ map[k3] = 30;
+ EXPECT_EQ(map.size(), 2U);
+ EXPECT_EQ(map.at(k1), 20);
+ EXPECT_EQ(map.at(k2), 20);
+ EXPECT_EQ(map.at(k3), 30);
+}
+
+TEST(StructuralKey, DefaultCtorInitializesHash) {
+ StructuralKeyObj obj;
+ EXPECT_EQ(obj.hash_i64, 0);
+}
+
+TEST(StructuralKey, MapAnyKeyUsesStructuralAttrs) {
+ Map<Any, Any> map;
+ StructuralKey k1(Array<int>{1, 2, 3});
+ StructuralKey k2(Array<int>{1, 2, 3});
+ StructuralKey k3(Array<int>{1, 3, 5});
+
+ map.Set(k1, 1);
+ map.Set(k2, 2);
+ map.Set(k3, 3);
+
+ // k1 and k2 are structurally equal and should occupy the same map key.
+ EXPECT_EQ(map.size(), 2U);
+ EXPECT_EQ(map.at(k1).cast<int>(), 2);
+ EXPECT_EQ(map.at(k2).cast<int>(), 2);
+ EXPECT_EQ(map.at(k3).cast<int>(), 3);
+}
+
+} // namespace
diff --git a/tests/python/test_structural.py b/tests/python/test_structural.py
new file mode 100644
index 00000000..df800f7c
--- /dev/null
+++ b/tests/python/test_structural.py
@@ -0,0 +1,92 @@
+# 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 tvm_ffi
+
+
+def test_structural_key_basic() -> None:
+ k1 = tvm_ffi.StructuralKey({"a": [1, 2], "b": [3, {"c": 4}]})
+ k2 = tvm_ffi.StructuralKey({"b": [3, {"c": 4}], "a": [1, 2]})
+ k3 = tvm_ffi.StructuralKey({"a": [1, 2], "b": [3, {"c": 5}]})
+
+ assert tvm_ffi.structural_hash(k1.key) == k1.__hash__()
+ assert tvm_ffi.structural_hash(k2.key) == k2.__hash__()
+
+ assert k1 == k2
+ assert k1 != k3
+ assert hash(k1) == hash(k2)
+ assert tvm_ffi.structural_equal(k1.key, k2.key)
+ assert not tvm_ffi.structural_equal(k1.key, k3.key)
+
+
+def test_structural_helpers() -> None:
+ lhs = {"items": [1, 2, {"k": 3}], "meta": {"tag": "x"}}
+ rhs = {"meta": {"tag": "x"}, "items": [1, 2, {"k": 3}]}
+ other = {"items": [1, 2, {"k": 4}], "meta": {"tag": "x"}}
+
+ assert tvm_ffi.structural_equal(lhs, rhs)
+ assert not tvm_ffi.structural_equal(lhs, other)
+ assert tvm_ffi.structural_hash(lhs) == tvm_ffi.structural_hash(rhs)
+ assert tvm_ffi.structural_hash(lhs) != tvm_ffi.structural_hash(other)
+ assert tvm_ffi.get_first_structural_mismatch(lhs, rhs) is None
+ assert tvm_ffi.get_first_structural_mismatch(lhs, other) is not None
+
+
+def test_structural_key_in_map() -> None:
+ k1 = tvm_ffi.StructuralKey({"x": [1, 2], "y": [3]})
+ k2 = tvm_ffi.StructuralKey({"y": [3], "x": [1, 2]})
+ k3 = tvm_ffi.StructuralKey({"x": [1, 2], "y": [5]})
+
+ m = tvm_ffi.Map({k1: 1, k2: 2, k3: 3})
+ assert len(m) == 2
+ assert m[k1] == 2
+ assert m[k2] == 2
+ assert m[k3] == 3
+
+
+def test_structural_key_in_python_dict() -> None:
+ k1 = tvm_ffi.StructuralKey({"name": ["a", "b"], "ver": [1]})
+ k2 = tvm_ffi.StructuralKey({"ver": [1], "name": ["a", "b"]})
+ k3 = tvm_ffi.StructuralKey({"name": ["a", "c"], "ver": [1]})
+
+ data = {k1: "a", k3: "b"}
+ assert data[k2] == "a"
+ assert data[k3] == "b"
+
+
+def test_structural_key_tensor_content_policy() -> None:
+ t1_np = np.array([1.0, 2.0, 3.0], dtype="float32")
+ t2_np = np.array([1.0, 2.0, 4.0], dtype="float32")
+ if not hasattr(t1_np, "__dlpack__"):
+ return
+
+ t1 = tvm_ffi.from_dlpack(t1_np)
+ t2 = tvm_ffi.from_dlpack(t2_np)
+
+ # Default policy compares tensor content.
+ assert not tvm_ffi.structural_equal(t1, t2)
+ # Optional policy can ignore tensor content.
+ assert tvm_ffi.structural_equal(t1, t2, skip_tensor_content=True)
+
+ # StructuralKey should follow default structural policy.
+ k1 = tvm_ffi.StructuralKey(t1)
+ k2 = tvm_ffi.StructuralKey(t2)
+ assert k1 != k2
+
+ data = {k1: "a", k2: "b"}
+ assert len(data) == 2
