make_jaxpr

class braincore.transform.make_jaxpr(fun, static_argnums=(), axis_env=None, return_shape=False, abstracted_axes=None, state_returns=('read', 'write'))

Creates a function that produces its jaxpr given example args.

Parameters:

• fun (Callable) – 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 (Union[int, Iterable[int]]) – See the jax.jit() docstring.

• axis_env (Optional[Sequence[tuple[Hashable, int]]]) – 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 jax.pmap().

• return_shape (bool) – 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[Any]) – 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 (Union[str, Tuple[str, ...]]) – 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.

Return type:

Callable[..., Union[Tuple[ClosedJaxpr, Tuple[State, ...]], Tuple[ClosedJaxpr, Tuple[State, ...], Any]]]

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. 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 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,) }