test: update tests for braodcast

metaopt · Oct 6, 2023 · 438ef1b · 438ef1b
1 parent e17276f
commit 438ef1b
Showing 1 changed file with 32 additions and 0 deletions.
diff --git a/tests/test_prefix_errors.py b/tests/test_prefix_errors.py
@@ -81,12 +81,20 @@ def test_different_types():
     lhs_treespec, rhs_treespec = optree.tree_structure(lhs), optree.tree_structure(rhs)
     optree.tree_map_(lambda x, y: None, lhs, rhs)
     assert lhs_treespec.is_prefix(rhs_treespec)
+    assert (
+        len(lhs_treespec.broadcast_to(rhs_treespec, range(lhs_treespec.num_leaves)))
+        == rhs_treespec.num_leaves
+    )
     () = optree.prefix_errors(lhs, rhs)
 
     lhs, rhs = {'a': 1, 'b': 2}, defaultdict(int, {'a': 1, 'b': [2, 3]})
     lhs_treespec, rhs_treespec = optree.tree_structure(lhs), optree.tree_structure(rhs)
     optree.tree_map_(lambda x, y: None, lhs, rhs)
     assert lhs_treespec.is_prefix(rhs_treespec)
+    assert (
+        len(lhs_treespec.broadcast_to(rhs_treespec, range(lhs_treespec.num_leaves)))
+        == rhs_treespec.num_leaves
+    )
     () = optree.prefix_errors(lhs, rhs)
 
 
@@ -266,30 +274,50 @@ def test_different_metadata():
     lhs_treespec, rhs_treespec = optree.tree_structure(lhs), optree.tree_structure(rhs)
     optree.tree_map_(lambda x, y: None, lhs, rhs)  # ignore key ordering
     assert lhs_treespec.is_prefix(rhs_treespec)
+    assert (
+        len(lhs_treespec.broadcast_to(rhs_treespec, range(lhs_treespec.num_leaves)))
+        == rhs_treespec.num_leaves
+    )
     () = optree.prefix_errors(lhs, rhs)
 
     lhs, rhs = defaultdict(list, {'a': 1, 'b': 2}), defaultdict(set, {'b': [4, 5], 'a': 3})
     lhs_treespec, rhs_treespec = optree.tree_structure(lhs), optree.tree_structure(rhs)
     optree.tree_map_(lambda x, y: None, lhs, rhs)  # ignore default factory
     assert lhs_treespec.is_prefix(rhs_treespec)
+    assert (
+        len(lhs_treespec.broadcast_to(rhs_treespec, range(lhs_treespec.num_leaves)))
+        == rhs_treespec.num_leaves
+    )
     () = optree.prefix_errors(lhs, rhs)
 
     lhs, rhs = {'a': 1, 'b': 2}, defaultdict(list, {'b': [4, 5], 'a': 3})
     lhs_treespec, rhs_treespec = optree.tree_structure(lhs), optree.tree_structure(rhs)
     optree.tree_map_(lambda x, y: None, lhs, rhs)  # ignore dictionary types
     assert lhs_treespec.is_prefix(rhs_treespec)
+    assert (
+        len(lhs_treespec.broadcast_to(rhs_treespec, range(lhs_treespec.num_leaves)))
+        == rhs_treespec.num_leaves
+    )
     () = optree.prefix_errors(lhs, rhs)
 
     lhs, rhs = OrderedDict({'b': 5, 'a': 4}), {'a': 1, 'b': [2, 3]}
     lhs_treespec, rhs_treespec = optree.tree_structure(lhs), optree.tree_structure(rhs)
     optree.tree_map_(lambda x, y: None, lhs, rhs)  # ignore dictionary types
     assert lhs_treespec.is_prefix(rhs_treespec)
+    assert (
+        len(lhs_treespec.broadcast_to(rhs_treespec, range(lhs_treespec.num_leaves)))
+        == rhs_treespec.num_leaves
+    )
     () = optree.prefix_errors(lhs, rhs)
 
     lhs, rhs = deque([1, 2], maxlen=None), deque([3, [4, 5]], maxlen=3)
     lhs_treespec, rhs_treespec = optree.tree_structure(lhs), optree.tree_structure(rhs)
     optree.tree_map_(lambda x, y: None, lhs, rhs)  # ignore maxlen
     assert lhs_treespec.is_prefix(rhs_treespec)
+    assert (
+        len(lhs_treespec.broadcast_to(rhs_treespec, range(lhs_treespec.num_leaves)))
+        == rhs_treespec.num_leaves
+    )
     () = optree.prefix_errors(lhs, rhs)
 
     lhs, rhs = FlatCache([None, 1]), FlatCache(1)
@@ -432,6 +460,10 @@ def test_no_errors():
         optree.tree_map_(lambda x, y: None, lhs, rhs)
         lhs_treespec, rhs_treespec = optree.tree_structure(lhs), optree.tree_structure(rhs)
         assert lhs_treespec.is_prefix(rhs_treespec)
+        assert (
+            len(lhs_treespec.broadcast_to(rhs_treespec, range(lhs_treespec.num_leaves)))
+            == rhs_treespec.num_leaves
+        )
         () = optree.prefix_errors(lhs, rhs)