Recipe for automatically plumbing PRNG keys

[prasanthcakewalk]

Hi all. I wanted to share a recipe for automatically plumbing PRNG keys, while still maintaining support for JAX transforms and protecting against inadvertent reuse of keys.

The basic idea is to manage a stateful key vendor using a context variable in the global namespace. Context variables can be used to place local (traced) objects into the global namespace and access them directly from further up the call stack without plumbing those objects in and out of functions:

from contextvars import ContextVar
x_cvar = ContextVar("x_cvar")

def foo():
    # This is an impure function 
    x_val = x_cvar.get()
    return x_val**2

def bar(x):
    # This is a pure function
    # The result of x_cvar.get() within bar() is fully determined by the input x
    with x_cvar.set(x):
        return bar()

The recipe builds on this idea, adding a bell and a whistle to protect against unsafe jit. The full recipe is described here as part of a short tutorial: https://prasanth.codeberg.page/treepie/features/jit-aware-contextvars.html#auto-plumbing-random-keys.

---

Some context: I've been working on a library that provides will provide a few simple abstract features to simplify JAX programming, reduce boiler plate, etc. It's inspired by ideas from all around the jax ecosystem. Looking through the documentation for Ref objects and JEP 28845 (on stateful randomness), I figured treepie's JitAwareContextVar feature might be of interest here, especially to @jakevdp (I really appreciate your helpful answers to various questions here and elsewhere, btw). Also, because of how framework-agnostic this particular feature is, I think it could be suitable to be absorbed into JAX itself.

Treepie's codebase and documentation still need work (there are some assorted bugs and one important feature to implement). But the implementation of JitAwareContextVar is contained within just these few files:

• https://codeberg.org/prasanth/treepie/src/branch/main/src/treepie/pie_utils/_jit_aware_contextvars.py
• https://codeberg.org/prasanth/treepie/src/branch/main/src/treepie/_internal_utils/_trace_stack_inspection.py
• https://codeberg.org/prasanth/treepie/src/branch/main/src/treepie/_internal_utils/_safe_context_manager.py
  The notebook corresponding to the tutorial can be found here:
• https://codeberg.org/prasanth/treepie/src/branch/main/docs/source/features/jit-aware-contextvars.ipynb

Thoughts and feedback appreciated.

[jakevdp]

Very interesting, thanks for sharing! I do worry a bit about pushing traced values into a global context – that may lead to issues.

Relatedly, have you seen the experimental stateful_rng? This might already do more-or-less what you have in mind: https://docs.jax.dev/en/latest/_autosummary/jax.experimental.random.stateful_rng.html

[prasanthcakewalk]

Thanks for getting back to me! I’d seen the experimental stateful rng. The context variable based construction has a few advantages over it:

1. Let’s say we’ve built a complicated function that uses stateful_rng and makes calls to other functions that use the stateful_rng (and so on). And let’s say we’ve done all this without manually plumbing the stateful_rng. Now, if we want to vmap the overall function, we’re out of luck. We’ll have to reimplement every function involved (not just the outer most one), plumbing the rng. This is somewhat of a deterrent, since even simple functions might eventually get used in a more complicated function that needs to be vmapped.
   With the context variable based construction, the keys are getting plumbed properly, just not manually. So vmap works, as does every other jax transform.
2. In the context variable based approach, the compiler is not restricted by the sequential dependence. The counter is an int (even within a jit-trace), so the compiler will see it as a closed over constant.

Re: pushing traced values into a global context: Stdlib’s ContextVar makes strong guarantees about values not bleeding out of the context they were defined in, as long as it is used properly. In Treepie’s implementation, several constraints are added to ensure that the context variables can only be used properly. So, even through traced values get placed in the global context, they’re fairly secure imo.