Use while loop from Nx appropriately

I’ve been trying to implement an algorithm for inverse transform sampling to generate samples from a binomial distribution, but it seems that it get into an infinity loop while using using while. This is my code:

binom = Nx.pow((1 - p), n)
      s = binom
      pq = p / (1 - p)

      {result, _, _, _, _, _} = 
        while {x = 0, binom, u, s, n, pq}, Nx.greater(u, s) do
          x = x + 1
          binom = binom * pq * (n - x) / (x + 1)
          s = s + binom
          {x, binom, u, s, n, p}
        end
      result

u [0, 1] is a random number generated by a uniform distribution. You have to increment s (which is the cumulative distribution) till it gets bigger than u. Can someone help me with it?

I think you’re facing an issue I’ve faced before in livebook (cc @josevalim).
Your code is probably failing to compile because while you do pass pq in the initial state of the while, you’re using p in the end.

This fails to compile in Nx depending on how p is defined.

However, Livebook sometimes will just loop infinitely instead of crashing when facing a compile error.

1 Like

Failed to compile. p, n, u are not defined.