[Feature request or Q&A about DrJit] Is there something like np.random.choice(input, size, p) in DrJit?

[Mephisto405]

Hello,

I want to sample some indices according to the probability mass function of the indices.
In the context of NumPy, np.random.choice(input, size, p) method does the job well.
I wonder if there is an analogous in DrJit to the NumPy function.

[Mephisto405]

Oh, and I upload this question on Mitsuba's discussion board since the DrJit repo does not have one. Sorry for that.

[njroussel]

We don't have a one-liner like numpy, but Mitsuba does have explicit distribution classes.
You can use it like this:

import mitsuba as mi
import drjit as dr
mi.set_variant('cuda_ad_rgb')

distr = mi.DiscreteDistribution([0.3, 0.6, 0.1])

input = mi.Float([1, 2, 3])
samples = mi.PCG32(size=10).next_float32()

out = dr.gather(mi.Float, input, distr.sample(samples))
print(out)