def shard_spec_nothing(tree: PyTree[Any]) -> ShardingSpec:
    """
    Creates a sharding specification where no sharding is performed.

    This effectively clones the tree structure but replaces every leaf with SpecReplicate.

    Args:
        tree: The input PyTree structure.

    Returns:
        A new PyTree matching the input structure, containing strictly SpecReplicate for all leaves.
    """

    return pytree.tree_map(lambda _: SpecReplicate(), tree, is_leaf=lambda x: isinstance(x, (torch.Tensor, list)))

def shard_spec_on_dim(tree: PyTree[Any], dim: int) -> ShardingSpec:
    """
    Creates a sharding specification to split all tensors in the tree on a specific dimension.

    Iterates over the input tree:
    *   If a leaf is a Tensor with enough dimensions, it is mapped to a SpecShard(dim) object.
    *   If a leaf is a list, it is mapped to a SpecShard(0) object (only dim=0 is allowed for lists).
    *   Other types and 0-dim tensors are mapped to SpecReplicate.

    Args:
        tree: The input PyTree structure.
        dim: The dimension index to shard eligible tensors on.

    Returns:
        A new PyTree matching the input structure, containing SpecShard or SpecReplicate objects.

    Raises:
        ValueError: If a tensor exists in the tree with rank less than or equal to 'dim'.
    """

    return pytree.tree_map(
        lambda x: _tree_item_to_shard(x, dim),
        tree,
        is_leaf=lambda x: isinstance(x, (torch.Tensor, list))
    )

def shard_tree(
        tree: TSameTree,
        sharding_spec: ShardingSpec,
        num_shards: int,
        enforce_even_split: bool
) -> tuple[TSameTree, ...]:
    """
    Shards a PyTree into a tuple of PyTrees, one for each shard rank.

    This function takes a single global data structure and splits it into `num_shards`
    structures.

    *   If a spec leaf is a ``SpecShard(dim)``, the tensor or list is split along that dimension,
        and the ``i``-th slice goes to the ``i``-th output tree.
    *   If a spec leaf is ``SpecReplicate``, the item is replicated (reference copy) to all
        output trees.

    Args:
        tree: The structure containing tensors to be sharded.
        sharding_spec: A structure matching 'tree' containing ``SpecShard`` or ``SpecReplicate`` objects.
        num_shards: The total number of shards to split the tensors into.
        enforce_even_split: If True, raises a ValueError if a tensor's dimension
            size is not perfectly divisible by ``num_shards``.

    Returns:
        A tuple of length ``num_shards``. Each element is a PyTree matching
        the structure of the input ``tree``, containing the local data for
        that specific rank.

    Raises:
        ValueError: If tree structures do not match, or valid sharding conditions
            are not met.
    """
    flat_spec, spec_struct = pytree.tree_flatten(sharding_spec)

    try:
        flat_tree = spec_struct.flatten_up_to(tree)
    except (ValueError, TypeError) as e:
        raise ValueError("Tree structure does not match sharding spec") from e

    sharded_leaves_per_node = [
        _shard_leaf_to_list(item, spec, num_shards, enforce_even_split)
        for item, spec in zip(flat_tree, flat_spec, strict=True)
    ]

    rank_leaves = list(zip(*sharded_leaves_per_node, strict=True))

    return tuple(spec_struct.unflatten(leaves) for leaves in rank_leaves)

def unshard_tree(
        sharded_trees: Sequence[TSameTree],
        sharding_spec: ShardingSpec
) -> TSameTree:
    """
    Combines a sequence of PyTrees (one per rank) into a single global PyTree.

    This is the inverse of ``shard_tree``. It iterates over the provided trees,
    gathering corresponding leaves from each rank.

    *   If the spec for a leaf is ``SpecShard(dim)``, the tensors from all ranks are
        concatenated (or stacked if ``do_stack=True``) along that dimension.
    *   If the spec is ``SpecReplicate``, it assumes the data is replicated
        and takes the item from the first rank.

    Args:
        sharded_trees: A sequence (list or tuple) of PyTrees. There must be
            one tree for each shard rank, and they must all share the same
            structure as ``sharding_spec``.
        sharding_spec: A structure matching the input trees containing
            ``SpecShard`` or ``SpecReplicate`` objects.

    Returns:
        A single PyTree where distinct shards have been merged into full tensors.

    Raises:
        ValueError: If ``sharded_trees`` is empty, or if unit structures do
            not match the spec.
    """
    if not sharded_trees:
        raise ValueError("sharded_trees sequence cannot be empty")

    flat_spec, spec_struct = pytree.tree_flatten(sharding_spec)

    flat_shards_per_rank = []
    for i, tree in enumerate(sharded_trees):
        try:
            leaves = spec_struct.flatten_up_to(tree)
        except (ValueError, TypeError) as e:
            raise ValueError(
                f"Structure mismatch at shard {i}: tree does not match sharding spec structure"
            ) from e

        flat_shards_per_rank.append(leaves)

    grouped_leaves = list(zip(*flat_shards_per_rank, strict=True))

    reconstructed_leaves = [
        _unshard_leaf_from_group(group, spec)
        for group, spec in zip(grouped_leaves, flat_spec, strict=True)
    ]

    return spec_struct.unflatten(reconstructed_leaves)