nn.stochastic context does not extend into jitted functions

[adarob]

This causes silent failures where nn.make_rng() inside of the jitted function always produces the same "random" numbers. This issue should be well-documented and errors should be thrown to protect against it.

[AlexeyG]

@adarob could you provide a minimal repro for this?

[adarob]

Doesn't work:

@jax.jit
def rnd():
  return (jax.random.randint(nn.make_rng(), (5,), 0, 10), 
          jax.random.randint(nn.make_rng(), (5,), 0, 10))

with nn.stochastic(jax.random.PRNGKey(0)):
  for _ in range(5):
    print(rnd())

Output:

(DeviceArray([0, 2, 0, 3, 2], dtype=int32), DeviceArray([4, 1, 7, 8, 9], dtype=int32))
(DeviceArray([0, 2, 0, 3, 2], dtype=int32), DeviceArray([4, 1, 7, 8, 9], dtype=int32))
(DeviceArray([0, 2, 0, 3, 2], dtype=int32), DeviceArray([4, 1, 7, 8, 9], dtype=int32))
(DeviceArray([0, 2, 0, 3, 2], dtype=int32), DeviceArray([4, 1, 7, 8, 9], dtype=int32))
(DeviceArray([0, 2, 0, 3, 2], dtype=int32), DeviceArray([4, 1, 7, 8, 9], dtype=int32))

Works

@jax.jit
def rnd(rng):
  with nn.stochastic(rng):
    return (jax.random.randint(nn.make_rng(), (5,), 0, 10),
            jax.random.randint(nn.make_rng(), (5,), 0, 10))

with nn.stochastic(jax.random.PRNGKey(0)):
  for _ in range(5):
    print(rnd(nn.make_rng()))

Output:

(DeviceArray([8, 5, 6, 6, 7], dtype=int32), DeviceArray([4, 9, 7, 1, 5], dtype=int32))
(DeviceArray([9, 3, 1, 6, 0], dtype=int32), DeviceArray([6, 0, 5, 3, 9], dtype=int32))
(DeviceArray([2, 7, 8, 8, 1], dtype=int32), DeviceArray([9, 2, 5, 0, 6], dtype=int32))
(DeviceArray([0, 1, 2, 8, 1], dtype=int32), DeviceArray([5, 4, 6, 1, 1], dtype=int32))
(DeviceArray([1, 8, 4, 8, 3], dtype=int32), DeviceArray([1, 3, 6, 6, 4], dtype=int32))

[levskaya]

In that second example you meant to write rng in place of nn.make_rng(), no?

[adarob]

I don't know which line you're referring to but it looks like what I intended.

[levskaya]

Ah, my apologies I misread it on the first read.

[jheek]

This is part of a larger issue concerning mixing states and jax transformations. nn.stochastic should throw an exception in this case because mixing jax transformations and internal state are ambigious. I will make a PR for this but it might lead to some false positives that need to be fixed.

[avital]

I think #125 is the PR that should address this.

Effectively, it should make your code @adarob throw an explicit error, and then you can decide how to deal with the PRNGs. E.g. if you're using vmap you will have to explicitly choose whether you split them or reuse the PRNG.

[jheek]

Btw we are also looking into automatically supporting things like stateful and stochastic in combination with jax transforms together with the Haiku folks and the jax core team. But for know we just try to avoid silent errors

[avital]

@jheek assigning to you because I believe you're looking into this

[jheek]

nn.stochastic correctly throws an error but it does now extend into init_by_shape (as of PR #159).