Function flatten_items

axlearn/common/utils.py:421–429

Flattens `tree` and returns a list of (path, value) pairs.

(
    tree: Nested[Tensor], separator: str = "/", is_leaf: Optional[Callable[[Any], bool]] = None
)

def flatten_items(
    tree: Nested[Tensor], separator: str = "/", is_leaf: Optional[Callable[[Any], bool]] = None
) -> Sequence[tuple[str, Tensor]]:
    """Flattens `tree` and returns a list of (path, value) pairs."""
    flat_paths_and_values, _ = jax.tree_util.tree_flatten_with_path(tree, is_leaf=is_leaf)
    return list(
        (separator.join(_key_entry_to_str(k) for k in path), value)
        for path, value in flat_paths_and_values
    )


@jax.tree_util.register_pytree_with_keys_class