Skip to contents

Sample from a discrete distribution, analogous to R's sample() function. Samples integers from 1 to n with uniform probability and with replacement.

Usage

nv_rdunif(shape, initial_state, n, dtype = "i32")

Arguments

shape

(integer())
Shape.

initial_state

([ui64[2]][tensorish])
RNG state.

n

(integer(1))
Number of categories to sample from (samples integers 1 to n).

dtype

(character(1) | stablehlo::TensorDataType)
Data type.

Value

(list() of tensorish)
List of two tensors: the new RNG state and the sampled integers.