import torch


def where(condition, x, y):
    """
    Return a tree of tensors of elements selected from either ``x`` or ``y``, depending on ``condition``.

    Examples::

        >>> import torch
        >>> import treetensor.torch as ttorch
        >>> ttorch.where(
        ...     torch.tensor([[True, False], [False, True]]),
        ...     torch.tensor([[2, 8], [16, 4]]),
        ...     torch.tensor([[3, 11], [5, 7]]),
        ... )
        tensor([[ 2, 11],
                [ 5,  4]])

        >>> tt1 = ttorch.randint(1, 99, {'a': (2, 3), 'b': {'x': (3, 2, 4)}})
        >>> tt1
        <Tensor 0x7f6760ad9908>
        ├── a --> tensor([[27, 90, 80],
        │                 [12, 59,  5]])
        └── b --> <Tensor 0x7f6760ad9860>
            └── x --> tensor([[[71, 52, 92, 79],
                               [48,  4, 13, 96]],

                              [[72, 89, 44, 62],
                               [32,  4, 29, 76]],

                              [[ 6,  3, 93, 89],
                               [44, 89, 85, 90]]])
        >>> ttorch.where(tt1 % 2 == 1, tt1, 0)
        <Tensor 0x7f6760ad9d30>
        ├── a --> tensor([[27,  0,  0],
        │                 [ 0, 59,  5]])
        └── b --> <Tensor 0x7f6760ad9f98>
            └── x --> tensor([[[71,  0,  0, 79],
                               [ 0,  0, 13,  0]],

                              [[ 0, 89,  0,  0],
                               [ 0,  0, 29,  0]],

                              [[ 0,  3, 93, 89],
                               [ 0, 89, 85,  0]]])
    """
    return torch.where(condition, x, y)