import torch


def ones_like(input, *args, **kwargs):
    """
    In ``treetensor``, you can use ``ones_like`` to create a tree of tensors with all ones like another tree.

    Example::

        >>> import torch
        >>> import treetensor.torch as ttorch
        >>> ttorch.ones_like(torch.randn(2, 3))  # the same as torch.ones_like(torch.randn(2, 3))
        tensor([[1., 1., 1.],
                [1., 1., 1.]])

        >>> ttorch.ones_like({
        ...     'a': torch.randn(2, 3),
        ...     'b': {'x': torch.randn(4, )},
        ... })
        <Tensor 0x7ff363bbc320>
        ├── a --> tensor([[1., 1., 1.],
        │                 [1., 1., 1.]])
        └── b --> <Tensor 0x7ff363bbc240>
            └── x --> tensor([1., 1., 1., 1.])
    """
    return torch.ones_like(input, *args, **kwargs)