def _tree_flatten(pytree: PyTree) -> tuple[list[Any], TreeSpec]:
    """Copy of :func:`torch.utils._pytree.tree_flatten` using our custom leaf function."""
    if _is_leaf_or_primitive_container(pytree):
        return [pytree], LeafSpec()

    node_type = _get_node_type(pytree)
    flatten_fn = SUPPORTED_NODES[node_type].flatten_fn
    child_pytrees, context = flatten_fn(pytree)

    result: list[Any] = []
    children_specs: list[TreeSpec] = []
    for child in child_pytrees:
        flat, child_spec = _tree_flatten(child)
        result += flat
        children_specs.append(child_spec)

    return result, TreeSpec(node_type, context, children_specs)
Copy of :func:`torch.utils._pytree.tree_flatten` using our custom leaf function.
src/lightning/pytorch/utilities/_pytree.py