# Copyright 2024 BDP Ecosystem Limited. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""
This module implements how to create a JAX Jaxpr from a given function by considering the states that are read and
written by the function. These state transformations are foundational for the BrainCore library. These utilities
include two basic functions: `StatefulFunction` and `make_jaxpr`.
``StatefulFunction``
--------------------
The module provides a class called ``StatefulFunction`` that wraps a function and provides methods to get the
JAX Jaxpr, the output shapes, the states that are read and written by the function, and the output of the function.
The class provides the following methods:
- `make_jaxpr`: creates the JAX Jaxpr of the function.
- `jaxpr_call`: calls the function at the JAX Jaxpr level.
- `jaxpr_call_without_states`: calls the function at the JAX Jaxpr level without considering the states.
- `get_states`: returns the states that are read and written by the function.
- `get_read_states`: returns the states that are read by the function.
- `get_write_states`: returns the states that are written by the function.
- `get_static_args`: returns the static arguments from the arguments.
- `compile_and_get_states_by_static_args`: compiles the function and returns the states that are read and
written by the function.
- `get_jaxpr`: returns the JAX Jaxpr of the function.
- `get_out_shapes`: returns the output shapes of the function.
- `get_out_treedef`: returns the output tree of the function.
``make_jaxpr``
--------------
The module provides a function called `make_jaxpr` that creates a function that produces its JAX Jaxpr given example
arguments. The function returns a wrapped version of the function that when applied to example arguments returns a
`ClosedJaxpr` representation of the function on those arguments. If the argument `return_shape` is `True`, then the
returned function instead returns a pair where the first element is the `ClosedJaxpr` representation of the function
and the second element is a pytree representing the structure, shape, dtypes, and named shapes of the output of the
function.
"""
from __future__ import annotations
import functools
import operator
from collections.abc import Hashable, Iterable, Sequence
from typing import Any, Callable, Tuple, Union, Dict, Optional
import jax
from jax._src import source_info_util
from jax.interpreters import partial_eval as pe
from jax.util import wraps
from jax.interpreters.xla import abstractify
from braincore._state import State, StateTrace
from braincore._common import set_module_as
PyTree = Any
__all__ = [
"StatefulFunction", "make_jaxpr",
]
def _assign_state_values(states, state_vals) -> None:
"""
Assign the state values to the states.
Args:
states: The states.
state_vals: The state values.
"""
assert len(states) == len(state_vals), f'State length mismatch. {len(states)} != {len(state_vals)}.'
for st, val in zip(states, state_vals):
st.value = val
def _ensure_index_tuple(x: Any) -> tuple[int, ...]:
"""Convert x to a tuple of indices."""
x = jax.core.concrete_or_error(None, x, "expected a static index or sequence of indices.")
try:
return (operator.index(x),)
except TypeError:
return tuple(jax.util.safe_map(operator.index, x))
def _new_arg(frame, trace, aval):
"""
Transform a new argument to a tracer.
Modified from jax.interpreters.partial_eval.DynamicJaxprTrace.new_arg()
Args:
frame: The frame.
trace: The trace.
aval: The abstract value.
Returns:
The tracer.
"""
tracer = pe.DynamicJaxprTracer(trace, aval, source_info_util.current())
frame.tracers.append(tracer)
frame.tracer_to_var[id(tracer)] = var = frame.newvar(aval)
frame.invars.append(var)
return tracer
def wrapped_abstractify(x: Any) -> Any:
"""
Abstractify the input.
Args:
x: The input.
Returns:
The abstractified input.
"""
if isinstance(x, pe.DynamicJaxprTracer):
return jax.core.ShapedArray(x.aval.shape, x.aval.dtype, weak_type=x.aval.weak_type)
return abstractify(x)
class StatefulFunction(object):
"""
A wrapper class for a function that collects the states that are read and written by the function. The states are
collected by the function and returned as a StateDictManager instance. The StateDictManager instance can be used to
manage the states in the JAX program. The class provides a function called `states` that returns the states
that are read and written by the function. The class provides a function called `to_state_manager` that returns
a StateDictManager instance that contains the states that are read and written by the function. The class provides
a function called `__call__` that wraps the function and returns the states that are read and written by the
function and the output of the function.
Args:
fun: The function whose ``jaxpr`` is to be computed. Its positional
arguments and return value should be arrays, scalars, or standard Python
containers (tuple/list/dict) thereof.
static_argnums: See the :py:func:`jax.jit` docstring.
axis_env: Optional, a sequence of pairs where the first element is an axis
name and the second element is a positive integer representing the size of
the mapped axis with that name. This parameter is useful when lowering
functions that involve parallel communication collectives, and it
specifies the axis name/size environment that would be set up by
applications of :py:func:`jax.pmap`.
abstracted_axes: Optional, a pytree with the same structure as the input
arguments to ``fun``. The leaves of the pytree can be either None or a
dict with axis names as keys and integers as values. If the leaf is None,
then the corresponding axis is not abstracted. If the leaf is a dict, then
the corresponding axis is abstracted, and the dict specifies the axis name
and size. The abstracted axes are used to infer the input type of the
function. If None, then all axes are abstracted.
state_returns: Optional, a string or a tuple of strings. The default is
``('read', 'write')``. The strings specify the categories of states to be
returned by the wrapped function. The categories are ``'read'`` and
``'write'``. If the category is ``'read'``, then the wrapped function
returns the states that are read by the function. If the category is
``'write'``, then the wrapped function returns the states that are written
by the function. If the category is ``'read'`` and ``'write'``, then the
wrapped function returns both the read and write states.
"""
__module__ = "braincore.transform"
def __init__(
self,
fun: Callable,
static_argnums: Union[int, Iterable[int]] = (),
axis_env: Optional[Sequence[tuple[Hashable, int]]] = None,
abstracted_axes: Optional[Any] = None,
state_returns: Union[str, Tuple[str, ...]] = ('read', 'write'),
cache_type: Optional[str] = None,
):
# explicit parameters
self.fun = fun
self.static_argnums = _ensure_index_tuple(tuple() if static_argnums is None else static_argnums)
self.axis_env = axis_env
self.abstracted_axes = abstracted_axes
self.state_returns = tuple(state_returns) if isinstance(state_returns, (tuple, list)) else (state_returns,)
assert cache_type in [None, 'jit']
self.cache_type = cache_type
# implicit parameters
self._jaxpr: Dict[Any, jax.core.ClosedJaxpr] = dict()
self._out_shapes: Dict[Any, PyTree] = dict()
self._jaxpr_out_tree: Dict[Any, PyTree] = dict()
self._state_trace: Dict[Any, StateTrace] = dict()
def __repr__(self) -> str:
return (f"{self.__class__.__name__}({self.fun}, "
f"static_argnums={self.static_argnums}, "
f"axis_env={self.axis_env}, "
f"abstracted_axes={self.abstracted_axes}, "
f"state_returns={self.state_returns})")
[docs]
def get_jaxpr(self, cache_key: Hashable = ()) -> jax.core.ClosedJaxpr:
"""
Read the JAX Jaxpr representation of the function.
Args:
cache_key: The hashable key.
Returns:
The JAX Jaxpr representation of the function.
"""
if cache_key not in self._jaxpr:
raise ValueError(f"the function is not called with the static arguments: {cache_key}")
return self._jaxpr[cache_key]
[docs]
def get_out_shapes(self, cache_key: Hashable = ()) -> PyTree:
"""
Read the output shapes of the function.
Args:
cache_key: The hashable key.
Returns:
The output shapes of the function.
"""
if cache_key not in self._out_shapes:
raise ValueError(f"the function is not called with the static arguments: {cache_key}")
return self._out_shapes[cache_key]
[docs]
def get_out_treedef(self, cache_key: Hashable = ()) -> PyTree:
"""
Read the output tree of the function.
Args:
cache_key: The hashable key.
Returns:
The output tree of the function.
"""
if cache_key not in self._jaxpr_out_tree:
raise ValueError(f"the function is not called with the static arguments: {cache_key}")
return self._jaxpr_out_tree[cache_key]
[docs]
def get_states(self, cache_key: Hashable = ()) -> Tuple[State, ...]:
"""
Read the states that are read and written by the function.
Args:
cache_key: The hashable key.
Returns:
The states that are read and written by the function.
"""
if cache_key not in self._state_trace:
raise ValueError(f"the function is not called with the static arguments: {cache_key}")
return tuple(self._state_trace[cache_key].states)
[docs]
def get_read_states(self, cache_key: Hashable = ()) -> Tuple[State, ...]:
"""
Read the states that are read by the function.
Args:
cache_key: The hashable key.
Returns:
The states that are read by the function.
"""
_state_trace = self._state_trace[cache_key]
return tuple([st for st, ty in zip(_state_trace.states, _state_trace.types) if ty == 'read'])
[docs]
def get_write_states(self, cache_key: Hashable = ()) -> Tuple[State, ...]:
"""
Read the states that are written by the function.
Args:
cache_key: The hashable key.
Returns:
The states that are written by the function.
"""
state_trace = self._state_trace[cache_key]
return tuple([st for st, ty in zip(state_trace.states, state_trace.types) if ty == 'write'])
[docs]
def get_arg_cache_key(self, *args, **kwargs) -> Tuple:
"""
Get the static arguments from the arguments.
Args:
*args: The arguments to the function.
Returns:
The static arguments.
"""
if self.cache_type == 'jit':
static_args, dyn_args = [], []
for i, arg in enumerate(args):
if i in self.static_argnums:
static_args.append(arg)
else:
dyn_args.append(arg)
dyn_args = jax.tree.map(wrapped_abstractify, jax.tree.leaves(dyn_args))
dyn_kwargs = jax.tree.map(wrapped_abstractify, jax.tree.leaves(kwargs))
return tuple([tuple(static_args), tuple(dyn_args), tuple(dyn_kwargs)])
elif self.cache_type is None:
num_arg = len(args)
return tuple(args[i] for i in self.static_argnums if i < num_arg)
else:
raise ValueError(f"Invalid cache type: {self.cache_type}")
[docs]
def compile_and_get_states_by_static_args(self, *args, **kwargs) -> Tuple[State, ...]:
"""
Get the states that are read and written by the function.
Args:
*args: The arguments to the function.
**kwargs: The keyword arguments to the function.
Returns:
The states that are read and written by the function.
"""
cache_key = self.get_arg_cache_key(*args, **kwargs)
if cache_key not in self._state_trace:
self.make_jaxpr(*args, **kwargs)
return self.get_states(cache_key)
[docs]
def clear_cache(self) -> None:
"""
Clear the compilation cache.
"""
self._jaxpr.clear()
self._out_shapes.clear()
self._jaxpr_out_tree.clear()
self._state_trace.clear()
@staticmethod
def _init_trace_and_newarg() -> StateTrace:
# Should be within the calling of ``jax.make_jaxpr()``
state_trace: StateTrace = StateTrace()
main = jax.core.thread_local_state.trace_state.trace_stack.stack[-1]
frame = main.jaxpr_stack[-1]
trace = pe.DynamicJaxprTrace(main, jax.core.cur_sublevel())
state_trace.set_new_arg(functools.partial(_new_arg, frame, trace))
return state_trace
def _wrapped_fun_to_eval(self, cache_key, *args, **kwargs) -> Tuple[Any, Tuple[State, ...]]:
"""
Wrap the function and return the states that are read and written by the function and the output of the function.
Args:
*args: The arguments to the function.
**kwargs: The keyword arguments to the function.
Returns:
A tuple of the states that are read and written by the function and the output of the function.
"""
# state trace
_state_trace = self._init_trace_and_newarg()
self._state_trace[cache_key] = _state_trace
with _state_trace:
out = self.fun(*args, **kwargs)
state_values = _state_trace.collect_values('read', 'write')
_state_trace.recovery_original_values()
# return states is not allowed
# checking whether the states are returned
for leaf in jax.tree.leaves(out):
if isinstance(leaf, State):
leaf._raise_error_with_source_info(ValueError(f"State object is not allowed to be returned: {leaf}"))
return out, state_values
[docs]
def make_jaxpr(self, *args, **kwargs):
"""Creates a function that produces its jaxpr given example args.
A ``ClosedJaxpr`` representation of ``fun`` on those arguments. If the
argument ``return_shape`` is ``True``, then the returned function instead
returns a pair where the first element is the ``ClosedJaxpr``
representation of ``fun`` and the second element is a pytree representing
the structure, shape, dtypes, and named shapes of the output of ``fun``.
Args:
*args: The arguments to the function.
**kwargs: The keyword arguments to the function.
"""
# static args
cache_key = self.get_arg_cache_key(*args, **kwargs)
if cache_key not in self._state_trace:
try:
# jaxpr
jaxpr, (out_shapes, state_shapes) = jax.make_jaxpr(
functools.partial(self._wrapped_fun_to_eval, cache_key),
static_argnums=self.static_argnums,
axis_env=self.axis_env,
return_shape=True,
abstracted_axes=self.abstracted_axes
)(*args, **kwargs)
# returns
self._jaxpr_out_tree[cache_key] = jax.tree.structure((out_shapes, state_shapes))
self._out_shapes[cache_key] = (out_shapes, state_shapes)
self._jaxpr[cache_key] = jaxpr
except Exception as e:
try:
self._state_trace.pop(cache_key)
except KeyError:
pass
raise e
return self
[docs]
def jaxpr_call(self, state_vals, *args, **kwargs) -> Any:
"""
Call the function at the JAX Jaxpr level.
Args:
state_vals: The state values.
*args: The arguments to the function.
**kwargs: The keyword arguments to the function.
Returns:
State values and the function output.
"""
# state checking
cache_key = self.get_arg_cache_key(*args, **kwargs)
states = self.get_states(cache_key)
assert len(state_vals) == len(states), 'State length mismatch.'
# parameters
args = tuple(args[i] for i in range(len(args)) if i not in self.static_argnums)
args = jax.tree.flatten((args, kwargs, state_vals))[0]
# calling the function
closed_jaxpr = self.get_jaxpr(cache_key)
out_treedef = self.get_out_treedef(cache_key)
jaxpr_outs = jax.core.eval_jaxpr(closed_jaxpr.jaxpr, closed_jaxpr.consts, *args)
# output processing
out, new_state_vals = out_treedef.unflatten(jaxpr_outs)
assert len(new_state_vals) == len(state_vals), 'State length mismatch.'
return new_state_vals, out
[docs]
def jaxpr_call_auto(self, *args, **kwargs) -> Any:
"""
Call the function at the JAX Jaxpr level with automatic state management.
Args:
*args: The arguments to the function.
**kwargs: The keyword arguments to the function.
Returns:
The output of the function.
"""
cache_key = self.get_arg_cache_key(*args, **kwargs)
states = self.get_states(cache_key)
state_vals, out = self.jaxpr_call([st.value for st in states], *args, **kwargs)
for st, val in zip(states, state_vals):
st.value = val
return out
@set_module_as("braincore.transform")
def make_jaxpr(
fun: Callable,
static_argnums: Union[int, Iterable[int]] = (),
axis_env: Optional[Sequence[tuple[Hashable, int]]] = None,
return_shape: bool = False,
abstracted_axes: Optional[Any] = None,
state_returns: Union[str, Tuple[str, ...]] = ('read', 'write')
) -> Callable[..., (Tuple[jax.core.ClosedJaxpr, Tuple[State, ...]] |
Tuple[jax.core.ClosedJaxpr, Tuple[State, ...], PyTree])]:
"""Creates a function that produces its jaxpr given example args.
Args:
fun: The function whose ``jaxpr`` is to be computed. Its positional
arguments and return value should be arrays, scalars, or standard Python
containers (tuple/list/dict) thereof.
static_argnums: See the :py:func:`jax.jit` docstring.
axis_env: Optional, a sequence of pairs where the first element is an axis
name and the second element is a positive integer representing the size of
the mapped axis with that name. This parameter is useful when lowering
functions that involve parallel communication collectives, and it
specifies the axis name/size environment that would be set up by
applications of :py:func:`jax.pmap`.
return_shape: Optional boolean, defaults to ``False``. If ``True``, the
wrapped function returns a pair where the first element is the XLA
computation and the second element is a pytree with the same structure as
the output of ``fun`` and where the leaves are objects with ``shape``,
``dtype``, and ``named_shape`` attributes representing the corresponding
types of the output leaves.
abstracted_axes: Optional, a pytree with the same structure as the input
arguments to ``fun``. The leaves of the pytree can be either None or a
dict with axis names as keys and integers as values. If the leaf is None,
then the corresponding axis is not abstracted. If the leaf is a dict, then
the corresponding axis is abstracted, and the dict specifies the axis name
and size. The abstracted axes are used to infer the input type of the
function. If None, then all axes are abstracted.
state_returns: Optional, a string or a tuple of strings. The default is
``('read', 'write')``. The strings specify the categories of states to be
returned by the wrapped function. The categories are ``'read'`` and
``'write'``. If the category is ``'read'``, then the wrapped function
returns the states that are read by the function. If the category is
``'write'``, then the wrapped function returns the states that are written
by the function. If the category is ``'read'`` and ``'write'``, then the
wrapped function returns both the read and write states.
Returns:
A wrapped version of ``fun`` that when applied to example arguments returns
a ``ClosedJaxpr`` representation of ``fun`` on those arguments. If the
argument ``return_shape`` is ``True``, then the returned function instead
returns a pair where the first element is the ``ClosedJaxpr``
representation of ``fun`` and the second element is a pytree representing
the structure, shape, dtypes, and named shapes of the output of ``fun``.
A ``jaxpr`` is JAX's intermediate representation for program traces. The
``jaxpr`` language is based on the simply-typed first-order lambda calculus
with let-bindings. :py:func:`make_jaxpr` adapts a function to return its
``jaxpr``, which we can inspect to understand what JAX is doing internally.
The ``jaxpr`` returned is a trace of ``fun`` abstracted to
:py:class:`ShapedArray` level. Other levels of abstraction exist internally.
We do not describe the semantics of the ``jaxpr`` language in detail here, but
instead give a few examples.
>>> import jax
>>> import braincore as bc
>>>
>>> def f(x): return jax.numpy.sin(jax.numpy.cos(x))
>>> print(f(3.0))
-0.83602
>>> jaxpr, states = bc.transform.make_jaxpr(f)(3.0)
>>> jaxpr
{ lambda ; a:f32[]. let b:f32[] = cos a; c:f32[] = sin b in (c,) }
>>> jaxpr, states = bc.transform.make_jaxpr(jax.grad(f))(3.0)
>>> jaxpr
{ lambda ; a:f32[]. let
b:f32[] = cos a
c:f32[] = sin a
_:f32[] = sin b
d:f32[] = cos b
e:f32[] = mul 1.0 d
f:f32[] = neg e
g:f32[] = mul f c
in (g,) }
"""
stateful_fun = StatefulFunction(fun, static_argnums, axis_env, abstracted_axes, state_returns)
@wraps(fun)
def make_jaxpr_f(*args, **kwargs):
stateful_fun.make_jaxpr(*args, **kwargs)
cache_key = stateful_fun.get_arg_cache_key(*args, **kwargs)
if return_shape:
return (stateful_fun.get_jaxpr(cache_key),
stateful_fun.get_states(cache_key),
stateful_fun.get_out_shapes(cache_key)[0])
else:
return (stateful_fun.get_jaxpr(cache_key),
stateful_fun.get_states(cache_key))
# wrapped jaxpr builder function
make_jaxpr_f.__module__ = "braincore.transform"
if hasattr(fun, "__qualname__"):
make_jaxpr_f.__qualname__ = f"make_jaxpr({fun.__qualname__})"
if hasattr(fun, "__name__"):
make_jaxpr_f.__name__ = f"make_jaxpr({fun.__name__})"
return make_jaxpr_f
