Returns a shallow copy of the input tree with empty subtrees pruned. If a tree would be made empty by removal of its subtrees, it will also be pruned. This is a shallow copy because leaf nodes (non-dict values) are not deep-copied. Args: in_tree: the input tree to be pruned.
(in_tree: Nested[Tensor])
| 2066 | |
| 2067 | def find_cycles(tree: Nested) -> dict[str, KeyPath]: |
| 2068 | """Find a cycle in pytree `tree` if one exists. |
| 2069 | |
| 2070 | This function finds a descendant which has reference equality with one of its own |
| 2071 | ancestors, if one exists. |
| 2072 | |
| 2073 | Args: |
| 2074 | tree: The tree to find cycles in. |
| 2075 | |
| 2076 | Returns: |
| 2077 | If no cycle is found, an empty dict. |
| 2078 | If a cycle is found a dict with keys: |
| 2079 | * descendant: The KeyPath to the descendant. |
| 2080 | * ancestor: The KeyPath to the ancestor. |
| 2081 | """ |
| 2082 | |
| 2083 | def _find_cycles(tree: Nested, *, key_path: KeyPath, seen: list[int]) -> dict[str, KeyPath]: |
| 2084 | # DFS and check if path to root contains repeats. |