import torch


def randint(*args, **kwargs):
    """
    In ``treetensor``, you can use ``randint`` to create a tree of tensors with numbers
    in an integer range.

    Examples::

        >>> import torch
        >>> import treetensor.torch as ttorch
        >>> ttorch.randint(10, (2, 3))  # the same as torch.randint(10, (2, 3))
        tensor([[3, 4, 5],
                [4, 5, 5]])

        >>> ttorch.randint(10, {'a': (2, 3), 'b': {'x': (4, )}})
        <Tensor 0x7ff363bb6438>
        ├── a --> tensor([[5, 3, 7],
        │                 [8, 1, 8]])
        └── b --> <Tensor 0x7ff363bb6240>
            └── x --> tensor([8, 8, 2, 4])
    """
    return torch.randint(*args, **kwargs)