Full Module RNG Scope

[joeryjoery]

Hi, if I implement a Module that makes use of make_rng, e.g., Dropout, I now have to manually specify the rng when I use the Module's call function. Though I completely agree that the RNG streams should be split into distinct paths like it is now, it is quite annoying that I have to bookkeep this at every call to Module.apply.

import jax
import flax.linen as nn


rng = jax.random.PRNGKey(0)


class A(nn.Module):

    def __call__(self):
        x = self.make_rng('x')
        y = self.make_rng('y')

        return jax.random.uniform(x, ()) + jax.random.uniform(y, ())

myA = A()

# Currently
keys = jax.random.split(rng, 3)
scope = {'params', 'x', 'y'}  # I don't want to do this manually

# Desired functionality (or something like this)
# scope = myA.get_rng_scope()

rngs = dict(zip(scope, keys))

v = myA.init(rngs)
_ = myA.apply(v, rngs=rngs)

Is there already something like this implemented? Or could this be done in the future?

I'd also argue that flax should allow me to just pass in one global RNG key to either init or apply.

# Even better (just use one key)
_ = myA.apply(v, rngs=jax.random.PRNGKey(0))

This functionality would be great to have with the option to pass in the RNGS explicitly. This has no consequences for reproducibility since the rngs need to be split somewhere anyway.

At the moment the RNG design forces me to write work-arounds in my code which requires me to pass around a namespace that tracks all registered scopes. But this is quite error-prone and ideally this would be handled internally.

[cgarciae]

Hey @joeryjoery,

Is there already something like this implemented? Or could this be done in the future?

Maybe this could be implemented by iteratively running apply under eval_shape and inspecting the InvalidRngError exceptions to collect the keys. Could look something like this:

rngs = module.get_rng_keys(*args, **kwargs)

You'd sill have to pass the same inputs to apply except for rngs. Main downside is that is could be a bit slow for big modules as you are still tracing, but its possible indeed.

I'd also argue that flax should allow me to just pass in one global RNG key to either init or apply.

Its tricky to decide how this unique RNG would behave under lifted transforms, maybe we could give it a name like default and have make_rng use default if the current name is not available. Sadly this would make typos silent errors. Its a good topic to think about though.

[joeryjoery]

Hi thanks for your reply. I like the idea of a tracing solution, however, if I were to pass the Module around then I'd have to call the get_rng_keys method multiple times (possibly). So I would need to cache the keys somehow, which isn't nice since the Module arguments are often not hashable. So this kind of defeats the purpose.

I also don't like that I have to parse the error/ exception string, which might be refactored in flax in the future.

Anyways, this is what I got to work,

import flax.linen as nn
import flax.errors

import jax
import jax.numpy as jnp
import jax.random


class A(nn.Module):

    @nn.compact
    def __call__(self):
        x = self.make_rng('x')
        y = self.make_rng('y')

        dense = nn.Dense(1)
        return dense(jnp.ones(5)) * jax.random.uniform(x, ()) + jax.random.uniform(y, ())

    def get_rng_keys(self, *args, **kwargs) -> set[jax.Array]:
        keys = dict()
        while True:
            try:
                jax.eval_shape(self.init, keys, *args, **kwargs)
                
            except flax.errors.InvalidRngError as e:
                # Parse Message: `None needs PRNG for "x" (https:// ....)`
                key_missing = e.args[0].split(' ')[4].strip('"')
                
                keys[key_missing] = jax.random.PRNGKey(0)  # Dummy key
                
            else:  # No Exceptions
                return set(keys.keys())


myA = A()

scope = myA.get_rng_keys()
print(scope)
>> {'params', 'y', 'x'}

rng = jax.random.PRNGKey(0)
key_dict = {k: v for k, v in zip(scope, jax.random.split(rng, num=len(scope)))}
print(key_dict)
>> {'params': Array([2467461003,  428148500], dtype=uint32), 'y': Array([3186719485, 3840466878], dtype=uint32), 'x': Array([2562233961, 1946702221], dtype=uint32)}

v = myA.init(key_dict)
out = myA.apply(v, rngs=key_dict)  # `params` is not needed here in key_dict
print(out)
>> [-0.6150426]

[joeryjoery]

Wanted to add to this, during init could flax, as you suggest, use the default rng and build a rng scope containter similarly to how the Variable (params) container is built?

This would not require iterative re-tracing, since it simply reuses existing functionality.

So example usage would be like this (as an idea),

# Use `rng` as a default for all RNGs
variables, key_scope = myA.init(rng, return_rng_scope=True)  # default: return_rng_scope=False
print(key_scope)
>> {'params', 'x', 'y'}

# OR

# Equivalent to manual `rng` control:
variables = myA.init({'params': rng, 'x': rng, 'y': rng})