import torch


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

    Example::

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

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