oneflow.randint

oneflow.randint(low=0, high, size, *, dtype=None, generator=None, device=None, placement=None, sbp=None, requires_grad=False)Tensor

Returns a tensor filled with random integers generated uniformly between low (inclusive) and high (exclusive).

The shape of the tensor is defined by the variable argument size.

The interface is consistent with PyTorch. The documentation is referenced from: https://pytorch.org/docs/1.10/generated/torch.randint.html.

Parameters
  • low (int, optional) – Lowest integer to be drawn from the distribution. Default: 0.

  • high (int) – One above the highest integer to be drawn from the distribution.

  • size (tuple or oneflow.Size) – Defining the shape of the output tensor. Can be a variable number of arguments or a collection like a list or tuple or oneflow.Size.

Keyword Arguments
  • dtype (oneflow.dtype, optional) – The desired data type of returned tensor. Default: flow.int64.

  • generator (oneflow.Generator, optional) –

  • device (oneflow.device, optional) – The desired device of returned local tensor. If None, uses the current device.

  • placement (oneflow.placement, optional) – The desired device of returned global tensor. If None, will construct local tensor.

  • sbp (oneflow.sbp, optional) – The desired sbp of returned global tensor. It must be equal with the numbers of placement.

  • requires_grad (bool, optional) – If autograd should record operations on the returned tensor. Default: False.

For example:

>>> import oneflow as flow
>>> generator = flow.Generator()
>>> generator.manual_seed(0) 
<oneflow._oneflow_internal.Generator object at ...>
>>> y = flow.randint(0, 5, (3,3), generator=generator) # construct local tensor
>>> y
tensor([[2, 2, 3],
        [4, 3, 4],
        [2, 4, 2]], dtype=oneflow.int64)
>>> y.is_global
False
>>> placement = flow.placement("cpu", ranks=[0])
>>> y = flow.randint(0, 5, (3,3), generator=generator, placement=placement, sbp=flow.sbp.broadcast) # construct global tensor
>>> y.is_global
True