Checks that trees have the same structure. Args: *trees: A sequence of (at least 2) trees to assert equal structure between. Raises: ValueError: If ``trees`` does not contain at least 2 elements. AssertionError: If structures of any two trees are different.
(*trees: ArrayTree)
| 1372 | |
| 1373 | @_static_assertion |
| 1374 | def assert_trees_all_equal_structs(*trees: ArrayTree) -> None: |
| 1375 | """Checks that trees have the same structure. |
| 1376 | |
| 1377 | Args: |
| 1378 | *trees: A sequence of (at least 2) trees to assert equal structure between. |
| 1379 | |
| 1380 | Raises: |
| 1381 | ValueError: If ``trees`` does not contain at least 2 elements. |
| 1382 | AssertionError: If structures of any two trees are different. |
| 1383 | """ |
| 1384 | if len(trees) < 2: |
| 1385 | raise ValueError( |
| 1386 | "assert_trees_all_equal_structs on a single tree does not make sense. " |
| 1387 | "Maybe you wrote `assert_trees_all_equal_structs([a, b])` instead of " |
| 1388 | "`assert_trees_all_equal_structs(a, b)` ?") |
| 1389 | |
| 1390 | first_treedef = jax.tree_util.tree_structure(trees[0]) |
| 1391 | other_treedefs = (jax.tree_util.tree_structure(t) for t in trees[1:]) |
| 1392 | for i, treedef in enumerate(other_treedefs, start=1): |
| 1393 | if first_treedef != treedef: |
| 1394 | raise AssertionError( |
| 1395 | f"Error in tree structs equality check: trees 0 and {i} do not match," |
| 1396 | f"\n tree 0: {first_treedef}" |
| 1397 | f"\n tree {i}: {treedef}") |
| 1398 | |
| 1399 | |
| 1400 | @_static_assertion |
no outgoing calls
no test coverage detected
searching dependent graphs…