Distribution probabilities. Must sum to one on the last axis.
DType of the output in case this can't be inferred.
Whether to also return the log probability of sampled result. This is usually used for differentiating through stochastic variables, e.g. in reinforcement learning.
Shape to be sampled from each random distribution.