Random number generation inside a while loop

[LindoNkambule]

Hi all,

I am trying to generate random numbers inside jax.lax.while_loop but struggling with conceptualizing/figuring out how I would set this up in JAX. The (oversimplified) numpy equivalent of what I would like to achieve is something like this:

from numpy import random

while True:
    x = random.random()
    y = random.random()
    z = random.random()

    eval1 = ... # uses some values from outside the loop and the random ones
    eval2 = ... # uses some values from outside the loop and the random ones
    if eval1 <= eval2:
        break

Any help/pointers to relevant documentation would be appreciated. Thanks

[jakevdp]

It might look something like this:

import jax

def cond_fun(carry):
  _, val1, val2 = carry
  return val1 > val2

def body_fun(carry):
  key, _, _ = carry
  key, subkey1, subkey2 = jax.random.split(key, 3)
  val1 = jax.random.uniform(subkey1)
  val2 = jax.random.uniform(subkey2)
  return key, val1, val2

key, val1, val2 = jax.lax.while_loop(cond_fun, body_fun, (jax.random.PRNGKey(0), 1, 0))
print(val1, val2)
# 0.7184181 0.80332327

This is the equivalent of the following Python while loop:

key = jax.random.PRNGKey(0)
val1 = 1
val2 = 0
while val1 > val2:
  key, subkey1, subkey2 = jax.random.split(key, 3)
  val1 = jax.random.uniform(subkey1)
  val2 = jax.random.uniform(subkey2)
print(val1, val2)

[LindoNkambule]

Got it. Thank you so much!