Jit-aware context variables¶
Imports for this tutorial:
from dataclasses import dataclass
import jax
import jax.numpy as jnp
import numpy as np
from treepie.pie_utils import JitAwareContextVar, plumb_and_transform
For users familiar with the stdlib’s ContextVar, treepie’s JitAwareContextVar is a jit-aware spin on it. But no worries if you’re not familiar with the concept yet.
A primer on context variables (sans jit-awareness)¶
Context variables offer a way to set and access a context-local state. Let’s see what that means with some examples. We begin by creating a context variable called foo:
foo = JitAwareContextVar()
This object has two public methods: get and set_on_enter. Let’s use the get method:
foo.get()
---------------------------------------------------------------------------
LookupError Traceback (most recent call last)
Cell In[3], line 1
----> 1 foo.get()
File <hidden>/treepie/pie_utils/_jit_aware_contextvars.py:161, in JitAwareContextVar.get(self)
159 value, trace_where_value_was_set = self._wrapped_context_var.get()
160 except LookupError:
--> 161 raise LookupError(f"No value found for {self!r}.") from None
163 if _detect_jit_boundary(trace_where_value_was_set):
164 raise LookupError(
165 f"The value found for {self!r} was assigned behind a jit-transform "
166 "boundary. It cannot be accessed in the current context."
167 )
LookupError: No value found for <JitAwareContextVar at 0x75a5233827a0>.
This fails with an error because there is no value set for foo yet. Setting a value is done using the set_on_enter method, which returns a context manager. When entering the context, the value is assigned to the context variable within that context:
with foo.set_on_enter(1):
print(f"Within this context, {foo.get() = }")
Within this context, foo.get() = 1
When you exit a context, the value gets reset:
foo.get()
---------------------------------------------------------------------------
LookupError Traceback (most recent call last)
Cell In[5], line 1
----> 1 foo.get()
File <hidden>/treepie/pie_utils/_jit_aware_contextvars.py:161, in JitAwareContextVar.get(self)
159 value, trace_where_value_was_set = self._wrapped_context_var.get()
160 except LookupError:
--> 161 raise LookupError(f"No value found for {self!r}.") from None
163 if _detect_jit_boundary(trace_where_value_was_set):
164 raise LookupError(
165 f"The value found for {self!r} was assigned behind a jit-transform "
166 "boundary. It cannot be accessed in the current context."
167 )
LookupError: No value found for <JitAwareContextVar at 0x75a5233827a0>.
Contexts can be nested. When exiting a context, the value gets reset to what it was before entering the context:
with foo.set_on_enter(1):
assert foo.get() == 1
with foo.set_on_enter(2):
assert foo.get() == 2
assert foo.get() == 1
Context variables, if set to mutable objects, can be mutated within contexts:
with foo.set_on_enter([1]):
foo.get().append(2)
assert foo.get() == [1, 2]
with foo.set_on_enter({"a": 1}):
foo.get()["b"] = 2
assert foo.get() == {"a": 1, "b": 2}
# here foo reverts to [1, 2]
foo.get().append(3)
assert foo.get() == [1, 2, 3]
The value assigned to foo can depend on the variables in the previous context (including the previous value of foo itself). Reverting is to the previous value and not to a copy of the previous value:
with foo.set_on_enter([1]):
foo.get().append(2)
assert foo.get() == [1, 2]
with foo.set_on_enter(foo.get()):
# foo's value is unchanged when entering this context
# We continue editing its value here:
foo.get().append(3)
assert foo.get() == [1, 2, 3]
# Here foo doesn't revert to [1, 2].
foo.get().append(4)
assert foo.get() == [1, 2, 3, 4]
Similar to stdlib’s context variables, treepie’s jit-aware context variables are also thread- and async-aware. For details see the documentation for stdlib’s ContextVars.
Why are context variables useful for JAX programming?¶
Consider the following:
x_cvar = JitAwareContextVar()
def square(value_for_x):
with x_cvar.set_on_enter(value_for_x):
x_sq = x_cvar.get() ** 2 # Accessing global variable x_cvar here
return x_sq
# or equivalently
def square(value_for_x):
with x_cvar.set_on_enter(value_for_x):
return x_cvar.get() ** 2
assert square(5) == 25
Here, a global variable named x_cvar is accessed within the function square. Despite this, it is functionally pure. This is because the return value of x_cvar.get() is completely determined by the input value_for_x.
The fact that we were able to place a local object into the global namespace and access it from the global namespace within a specific context, all without affecting the behaviour outside of the specific context is powerful. The previous example is perhaps too simple to showcase just how powerful this is. Of course, square can be more straightforwardly defined as:
def simple_square(value_for_x):
return value_for_x**2
But now consider the following:
x_cvar = JitAwareContextVar()
def square_of_f(value_for_x):
with x_cvar.set_on_enter(value_for_x):
return impure_f() ** 2
def impure_f():
# Let f be jnp.log(x)
# Retrieve value for x:
x = x_cvar.get()
# Computational logic:
return jnp.log(x)
assert square_of_f(5.0) == jnp.log(5.0) ** 2
Here, clearly impure_f is not functionally pure. Its output depends on more than just its inputs (there are none in this case). However, square_of_f is still functionally pure. Let’s verify that:
@jax.jit
def jitted_square_of_f(value_for_x):
print("jit-tracing...")
return square_of_f(value_for_x)
for x_value in jnp.arange(1.0, 11.0):
assert jnp.isclose(jitted_square_of_f(x_value), jnp.log(x_value) ** 2)
print("Results match")
jit-tracing...
Results match
The function is jit-traced only once, yet its output is correct for different inputs, demonstrating functional purity.
The point to note here is that we were able to access value_for_x from within the function f without threading (or plumbing) it through as an input. This is why context variables are useful to JAX programming.
An aside
Context variables are a way to emulate algebraic effects, which is a modern concept in functional programming to declare and handle side-effects in a clean, type-safe, modular, and composable way. In this analogy, whatever functionality is implemented using the value retrieved via .get() is the effect. The .set_on_enter method is used to provide an effect handler. In actual algebraic effect systems, effects propagate down the call-stack until an appropriate handler is found. With context variables, retrieve requests propagate down the context-stack until an assigned value is found.
Everything described so far matches the ContextVar feature in Python’s stdlib (barring a few minor differences discussed near the end of this tutorial). So, what does “jit-awareness” add (or mean)?
Jit-awareness¶
The risk with the previous example is that one might inadvertently jit-trace the impure function inpure_f. Let’s see what happens if we do that:
# Sidenote: More verbose exception-traceback
# printing has been turned on for this code-cell.
x_cvar = JitAwareContextVar()
def square_of_f(value_for_x):
with x_cvar.set_on_enter(value_for_x):
return impure_f() ** 2
@jax.jit
def impure_f():
# Let f be jnp.log(x)
# Retrieve value for x:
x = x_cvar.get()
# Computational logic:
return jnp.log(x)
square_of_f(5.0)
---------------------------------------------------------------------------
LookupError Traceback (most recent call last)
Cell In[13], line 22
18 # Computational logic:
19 return jnp.log(x)
---> 22 square_of_f(5.0)
Cell In[13], line 9, in square_of_f(value_for_x)
7 def square_of_f(value_for_x):
8 with x_cvar.set_on_enter(value_for_x):
----> 9 return impure_f() ** 2
[... skipping hidden 5 frame]
Cell In[13], line 16, in impure_f()
12 @jax.jit
13 def impure_f():
14 # Let f be jnp.log(x)
15 # Retrieve value for x:
---> 16 x = x_cvar.get()
18 # Computational logic:
19 return jnp.log(x)
File <hidden>/treepie/pie_utils/_jit_aware_contextvars.py:164, in JitAwareContextVar.get(self)
161 raise LookupError(f"No value found for {self!r}.") from None
163 if _detect_jit_boundary(trace_where_value_was_set):
--> 164 raise LookupError(
165 f"The value found for {self!r} was assigned behind a jit-transform "
166 "boundary. It cannot be accessed in the current context."
167 )
169 return value
LookupError: The value found for <JitAwareContextVar at 0x75a52148e980> was assigned behind a jit-transform boundary. It cannot be accessed in the current context.
Jit-aware context variables protect you from performing such unsafe jit-traces by throwing an error, as demonstrated above. Here we’re using jit-trace to refer to any trace in preparation of a compilation, so the user is protected from unsafe jit.pmap and jit. shard_map transforms as well.
How it works
When a value is set using .set_on_enter, the inner-most JAX “trace” that is active at that time is recorded. Subsequently, when a value is accessed using .get, a walk is performed outwards, starting from the inner-most trace active at the time of the access. During the walk, if a jit-trace is detected between the trace where the value was set and the trace where the value was accessed (exclusive of the former, inclusive of the latter), a LookupError error is raised.
Trace semantics: Roughly speaking, when entering a JAX transform (jax.{jit, grad, vmap, pmap, shard_map, etc.}), corresponding trace object(s) is added to the stack of active traces. When the transform ends, the trace is removed from the stack. Entering tracing for jit, pmap, shard_map, etc. adds a specific kind of trace object (referred to internally in the jax-codebase as “DynamicJaxprTrace”) to the stack. Jit-, pmap-, shard_map-boundaries can be detected by walking down the stack and looking for this kind of trace object.
When to use context-variables: The rationale behind context variables is as follows: Let’s say you’re implementing some callable that is expected to primarily be a sub-component of a larger computatation. And let’s say some input to this sub-computation is “boring”; it will either be known way down the call-stack or will evolve in a mechanical manner before reaching the computation in question. For example, (i) a common setting value that is used by a many layers of a neural network (but you don’t want to mark it as static or close over it), or (ii) random keys evolving from some starting value in a mechanical way before reaching a given callable as input. In such cases, instead of plumbing these boring objects through the different function calls (possibly in a type-unsafe manner), you can just querying the context-stack for their value. You trust that some more significant encompassing computation will supply the values for these variables and call the sub-computation in the corresponding context. The jit-awareness protects you from accidentally jitting some callable without properly tracing over the context-variables used within it.
Before proceeding to some recipes, let’s look at one more, admitted bizarre, feature.
Plumb and transform¶
Let’s say you want to jit some sub-computation, even though it’s just a component of the actual computation of interest. For example, let’s say we want to jit the main computational logic in impure_f in the previous code-cell. However, we still want to be able to call impure_f from square_of_f without plumbing an input. This can be achieved as follows:
# Sidenote: More verbose exception-traceback
# printing has been turned on for this code-cell.
x_cvar = JitAwareContextVar()
def square_of_f(value_for_x):
with x_cvar.set_on_enter(value_for_x):
return impure_f() ** 2
def _orig_impure_f():
# Let f be jnp.log(x)
# Retrieve value for x:
x = x_cvar.get()
# Computational logic:
return jnp.log(x)
@jax.jit
def _jitted_pure_f(x):
print("jit-tracing...")
with x_cvar.set_on_enter(x):
return _orig_impure_f()
def impure_f():
return _jitted_pure_f(x_cvar.get())
for x_value in jnp.arange(1.0, 11.0):
assert jnp.isclose(square_of_f(x_value), jnp.log(x_value) ** 2)
print("Results match")
jit-tracing...
Results match
What’s happening here: We begin with an _orig_inpure_f, which depends on the context variable x_cvar. We want to jit the computational logic in this function. So we make a pure wrapper around it, which takes x as an input and provides it to _orig_impure_f. This function can now be jitted which leads to _jitted_pure_f. However, we want an impure version of the function to be available to the rest of the codebase, so we create a another wrapper impure_f, which calls _jitted_pure_f, using retrieving the value of x from the context variable.
If one buys into context variables, then this double-wrapper pattern will likely be used a lot. The plumb_and_transform function is a utility to automate this pattern:
“Yeah, but your scientists were so preoccupied with whether or not they could, they didn’t stop to think if they should.” - Dr. Ian Malcolm
x_cvar = JitAwareContextVar()
def square_of_f(value_for_x):
with x_cvar.set_on_enter(value_for_x):
return impure_f() ** 2
@plumb_and_transform(
transform_to_apply=jax.jit, # One can provide kwargs to jit here, or use any other transform.
compute_extra_input_before=lambda: x_cvar.get(),
create_context_managers_after=lambda x_value: [x_cvar.set_on_enter(x_value)],
)
def impure_f():
# Let f be jnp.log(x)
# Retrieve value for x:
x = x_cvar.get()
# Computational logic:
return jnp.log(x)
for x_value in jnp.arange(1.0, 11.0):
assert jnp.isclose(square_of_f(x_value), jnp.log(x_value) ** 2)
print("Results match")
Results match
# We can see that `impure_f` wraps around a jitted function here:
type(impure_f.__closure__[2].cell_contents)
jaxlib._jax.PjitFunction
Sidenotes: plumb_and_transform is designed to work regardless of the call signature of the the function being transformed. The API for plumb_and_transform is still being workshopped. Right now, the utility can plumb values into a wrapper function before transforming it. Eventually the utility will be able to plumb values out as well. This tutorial will be updated with more details after the API settles down.
Recipes¶
Auto-plumbing random keys¶
Let’s see how context variables can be used to construct functionally pure stochastic transformations without manually threading random keys in and out of various callables. We begin by creating a stateful KeyStream class:
from threading import Lock
# Jax key object (possibly a tracer)
type Key = jax.Array
# int or 0-dim immutable integer-like object (possibly a tracer)
type IntScalar = int | jax.Array | np.number
class KeyStream:
__slots__ = ("_base_key", "_offset", "_lock", "counter")
_base_key: Key
_offset: IntScalar
_lock: Lock
counter: int
def __init__(self, base_key: Key, offset: IntScalar = 0) -> None:
# Maybe do some input validation here
self._base_key = base_key
self._offset = offset
self._lock = Lock()
self.counter = 0
def next_key(self) -> Key:
with self._lock: # For thread safety
fold_in_data = self._offset + self.counter
self.counter += 1
return jax.random.fold_in(self._base_key, fold_in_data)
def __eq__(self, other) -> bool:
return self is other
__hash__ = None
Next we create a context variable to set and get KeyStream objects:
# Sidenote: `JitAwareContextVar` is a generic type,
# parameterized by the type of the objects being stored in it.
_keystream_contextvar: JitAwareContextVar[KeyStream] = JitAwareContextVar()
For convenience, let’s define some utility functions to manage/use the context variable:
def set_keystream_on_enter(keystream: KeyStream):
assert isinstance(keystream, KeyStream)
return _keystream_contextvar.set_on_enter(keystream)
def get_next_key() -> Key:
keystream = _keystream_contextvar.get()
return keystream.next_key()
Sidenote
One may want to put the definitions of KeyStream, _keystream_contextvar, set_keystream_on_enter and get_next_key in a separate module/file, to reduce the risk of _keystream_contextvar accidentally getting clobbered in the global namespace or being overridden by a local variable of the same name. This will also help standardize the usage of pseudo-randomness across a codebase.
Now let’s see the functions set_keystream_on_enter and get_next_key in action:
# Within stochastic functions and callables, we create random keys as needed, using get_next_key().
# We're not manually threading (or plumbing) the key through different calls.
# There's no "rng" or "key" or "**kwargs" among the function arguments.
def func_1() -> jax.Array:
return jax.random.normal(get_next_key())
def func_2(x: jax.Array) -> jax.Array:
return x + jax.random.normal(get_next_key()) + func_1()
def func_main(x: jax.Array):
return x * func_2(x) * func_1()
Tip
If one creates a wrapper around jax.random.normal as follows:
def normal(shape=(), dtype=None, *, out_sharding=None):
return jax.random.normal(key=get_next_key(), shape=shape, dtype=dtype, out_sharding=out_sharding)
and so on for other distributions, then the previous code-cell could be implemented without the word “key” in sight.
Having completed the definition of the computational graph (albeit functionally impurely, because of the get_next_key() calls), it can be made functionally pure and jitted as follows:
def pure_func_main(x: jax.Array, base_key: Key):
keystream = KeyStream(base_key)
with set_keystream_on_enter(keystream):
return func_main(x)
@jax.jit
def jitted_pure_func_main(x: jax.Array, base_key: Key):
print("jit-tracing...")
return pure_func_main(x, base_key)
# Calling the non-jitted function on multiple inputs
base_keys = jax.random.split(jax.random.key(42), 4)
x_values = jnp.arange(1.0, 5.0)
for x, base_key in zip(x_values, base_keys):
print(pure_func_main(x, base_key))
-0.1265627
0.41922864
-6.4566073
25.958435
# Calling the jitted version of the function on the same inputs
for x, base_key in zip(x_values, base_keys):
print(jitted_pure_func_main(x, base_key))
jit-tracing...
-0.12656265
0.41922864
-6.456607
25.958435
It works as intended.
Because of the functional purity of pure_func_main, jax.vmap also works as intended (output matches expectations; jit-traced only once for fixed shapes and types of input):
@jax.jit
@jax.vmap
def jitted_vmapped_pure_func_main(x: jax.Array, base_key: Key):
print("jit-tracing...")
return pure_func_main(x, base_key)
print(jitted_vmapped_pure_func_main(x_values[0:2], base_keys[0:2]))
print(jitted_vmapped_pure_func_main(x_values[2:4], base_keys[2:4]))
jit-tracing...
[-0.12656265 0.41922864]
[-6.456607 25.958435]
Sidenote
During jit-tracing, each get_next_key() call represents a sub-computation that depends only on (a) the base_key and a possible offset (neither of which are intended to change during the overall computation being traced), and (b) an int-valued counter, which is a closed-over constant (albeit, a different constant for each get_next_key() call). Consequently, during compilation, the compiler can reorder operations that depend on these keys without changing the output.
As before, one is protected from advertently jitting a function that doesn’t properly trace over base_key:
bad_main = jax.jit(func_main)
bad_main(jnp.array(1.0))
---------------------------------------------------------------------------
LookupError Traceback (most recent call last)
Cell In[25], line 3
1 bad_main = jax.jit(func_main)
----> 3 bad_main(jnp.array(1.0))
[... skipping hidden 5 frame]
Cell In[20], line 15, in func_main(x)
14 def func_main(x: jax.Array):
---> 15 return x * func_2(x) * func_1()
Cell In[20], line 11, in func_2(x)
10 def func_2(x: jax.Array) -> jax.Array:
---> 11 return x + jax.random.normal(get_next_key()) + func_1()
Cell In[19], line 7, in get_next_key()
6 def get_next_key() -> Key:
----> 7 keystream = _keystream_contextvar.get()
8 return keystream.next_key()
File <hidden>/treepie/pie_utils/_jit_aware_contextvars.py:161, in JitAwareContextVar.get(self)
159 value, trace_where_value_was_set = self._wrapped_context_var.get()
160 except LookupError:
--> 161 raise LookupError(f"No value found for {self!r}.") from None
163 if _detect_jit_boundary(trace_where_value_was_set):
164 raise LookupError(
165 f"The value found for {self!r} was assigned behind a jit-transform "
166 "boundary. It cannot be accessed in the current context."
167 )
LookupError: No value found for <JitAwareContextVar at 0x75a521466a20>.
with set_keystream_on_enter(KeyStream(jax.random.key(42))):
bad_main(jnp.array(1.0))
---------------------------------------------------------------------------
LookupError Traceback (most recent call last)
Cell In[26], line 2
1 with set_keystream_on_enter(KeyStream(jax.random.key(42))):
----> 2 bad_main(jnp.array(1.0))
[... skipping hidden 5 frame]
Cell In[20], line 15, in func_main(x)
14 def func_main(x: jax.Array):
---> 15 return x * func_2(x) * func_1()
Cell In[20], line 11, in func_2(x)
10 def func_2(x: jax.Array) -> jax.Array:
---> 11 return x + jax.random.normal(get_next_key()) + func_1()
Cell In[19], line 7, in get_next_key()
6 def get_next_key() -> Key:
----> 7 keystream = _keystream_contextvar.get()
8 return keystream.next_key()
File <hidden>/treepie/pie_utils/_jit_aware_contextvars.py:164, in JitAwareContextVar.get(self)
161 raise LookupError(f"No value found for {self!r}.") from None
163 if _detect_jit_boundary(trace_where_value_was_set):
--> 164 raise LookupError(
165 f"The value found for {self!r} was assigned behind a jit-transform "
166 "boundary. It cannot be accessed in the current context."
167 )
169 return value
LookupError: The value found for <JitAwareContextVar at 0x75a521466a20> was assigned behind a jit-transform boundary. It cannot be accessed in the current context.
However, one can create another util function to plumb and jit:
_MISSING = object()
# The depth of the key tree increases by 1 when entering a
# function decorated with `plumb_rng_key_and_jit`
def plumb_rng_key_and_jit(func=_MISSING, **jit_kwargs):
if func is _MISSING:
return lambda func: plumb_rng_key_and_jit(func, **jit_kwargs)
return plumb_and_transform(
callable_to_transform=func,
transform_to_apply=jax.jit(**jit_kwargs),
compute_extra_input_before=get_next_key,
create_context_managers_after=lambda key: [
set_keystream_on_enter(KeyStream(key))
],
)
bad_main = plumb_rng_key_and_jit(func_main)
with set_keystream_on_enter(KeyStream(jax.random.key(42))):
print(bad_main(jnp.array(1.0)))
-0.12656265
Throw in a default value for the key stream context variable, and in principle one can completely do way with manually key management, while still providing users an option to manually set keys in specific contexts.
Stateful nodes in pytrees¶
Another application of context variables is auto-plumbling state-objects into and out of methods of stateful nodes (e.g. stateful layers in NNs). Let’s say all the stateful nodes in a pytree have somehow been assigned a unique, hashable id that is invariant under tree operations like flatten/unflatten. Let state_dict represent a dict of the states of nodes, indexed by their ids. This state_dict can be automatically plumbed as follows:
# This is a basic sketch of the recipe
# Create a context variable
# This wuld likely be specific to a particular "model".
state_dict_contextvar: JitAwareContextVar[dict] = JitAwareContextVar()
@jax.jit
def main_func(tree, state_dict: dict):
with state_dict_contextvar.set_on_enter(state_dict):
# some computation
result = ... # some more computation
return (
result,
statedict_context_var.get(),
) # Return the state so it is a traced output
Methods of stateful nodes can access and update their state as follows:
# Registering SomeStatefulNode as a pytree node is not shown here
class SomeStatefulNode:
...
def some_method(self, some_params):
stateful_state = state_dict_contextvar.get()
cur_state = state_dict[self.unique_id]
# some computation
result = ... # some more computation
new_state = ... # some more computation
state_dict[unique_id] = new_state
return result
Alternatively, one can create a base class for stateful nodes as follows:
# Registering SomeStatefulNode as a pytree node is not shown here
class StatefulNodeBase:
def get_cur_state(self):
return state_dict_contextvar.get()[self.unique_id]
def update_state(self, new_state):
state_dict_contextvar.get()[self.unique_id] = new_state
class SomeStatefulNode(StatefulNodeBase):
...
def some_method(self, some_params):
cur_state = get_cur_state()
# some computation
result = ... # some more computation
new_state = ... # some more computation
self.update_state(new_state_arr)
return result
Tip
For fans of strong typing: If you want type inference and checking for state retrievals and updates, you can make StatefulNodeBase generic with a type parameter corresponding to the state:
class StatefulNodeBase[StateT]:
def get_cur_state(self) -> StateT:
return state_dict_contextvar.get()[self.unique_id]
def update_state(self, new_state: StateT) -> None:
state_dict_contextvar.get()[self.unique_id] = new_state
More specifications¶
Warning: context variables don’t get properly garbage collected.
Due to the implementation details of stdlib’s ContextVar, treepie’s JitAwareContextVar objects don’t get properly garbage collected. So, it is recommended that context variables (both stdlib’s and treepie’s) only be created at the top module level (possibly as a class-variable in a class defined at the top module level) and never in closures or loops. Context variables are usually created with the expectation that they’ll stay alive until the program ends.
To be clear, this is unrelated to the garbage collection of the values assigned to context variables using .set_on_enter(); this is about the garbage collection of the context variable itself.
Under construction
Spec and API differences from stdlib’s ContextVar (apart from jit-awareness)
Under construction
Limitations of context variables¶
Jit-awareness of
JitAwareContextVaris not a silver bullet to the issue of accidental side-effects. One still needs to ensure that the object being assigned to a context variable viaset_on_entercan be computed using just the traced inputs (and possible intended closures and static args) of an enclosing function call being jitted.…