crystal-data / num.cr

Scientific computing in pure Crystal
MIT License
151 stars 12 forks source link

feat: add mutlinomial sampling for tensors #88

Closed LucianBuzzo closed 1 year ago

LucianBuzzo commented 1 year ago

This change adss a multinomial method to the Tensor class, allowing you to draw samples from a multinomial distribution. The multinomial method accepts a tensor input and will produce an output tensor that samples the input probabilities.

Num::Rand.set_seed(0)
input = [[0.5, 0.5], [0.5, 0.5]].to_tensor
a = Tensor.multinomial(input, 5)
puts a # => [[0, 1, 1, 0, 1], [1, 0, 1, 1, 0]]

input2 = [0.5, 0.5, 0.5, 0.5].to_tensor
b = Tensor.multinomial(input, 6)
puts b # => [3, 2, 1, 1, 0, 2]

The logic of this method is based on the equivalent torch.multinomial method: https://pytorch.org/docs/stable/generated/torch.multinomial.html

christopherzimmerman commented 1 year ago

Looks great! Thanks for the PR!

LucianBuzzo commented 1 year ago

@christopherzimmerman Would you be able to cut a new release? Say v1.1.0? Otherwise I need to use a commit sha to get access to the new function - thanks! 🙏

christopherzimmerman commented 1 year ago

@LucianBuzzo just tagged a release