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