Returns a tree of the same structure as `nested_tensor` but with corresponding paths instead of values. E.g., tree_paths({'a': 1, 'b': [2, {'c': 3}]}) = {'a': 'a', 'b': ['b/0', {'c': 'b/1/c'}]} Args: tree: A nested structure. separator: The separator between par
(
tree: NestedTree, separator: str = "/", is_leaf: Optional[Callable[[Any], bool]] = None
)
| 394 | |
| 395 | |
| 396 | def tree_paths( |
| 397 | tree: NestedTree, separator: str = "/", is_leaf: Optional[Callable[[Any], bool]] = None |
| 398 | ) -> NestedTree: |
| 399 | """Returns a tree of the same structure as `nested_tensor` but with corresponding paths instead |
| 400 | of values. |
| 401 | |
| 402 | E.g., |
| 403 | tree_paths({'a': 1, 'b': [2, {'c': 3}]}) = {'a': 'a', 'b': ['b/0', {'c': 'b/1/c'}]} |
| 404 | |
| 405 | Args: |
| 406 | tree: A nested structure. |
| 407 | separator: The separator between parts of a path. |
| 408 | is_leaf: A Callable to evaluate whether the given node should be considered a leaf when |
| 409 | it otherwise would not, similarly to the is_leaf in jax.tree.map. |
| 410 | |
| 411 | Returns: |
| 412 | A nested structure with the same structure as `tree`, but each leaf will be a string path. |
| 413 | Note that None is not considered a leaf by jax.tree_util, hence also preserved by |
| 414 | tree_paths. |
| 415 | """ |
| 416 | return jax.tree.map_with_path( |
| 417 | lambda kp, _: separator.join(_key_entry_to_str(k) for k in kp), tree, is_leaf=is_leaf |
| 418 | ) |
| 419 | |
| 420 | |
| 421 | def flatten_items( |