brainstate.transform.eval_shape

Contents

brainstate.transform.eval_shape#

brainstate.transform.eval_shape(f, *args, return_state_shapes=False, **kwargs)[source]#

Compute the abstract output shape of f without executing it.

This is the brainstate-aware analogue of jax.eval_shape(). It traces f abstractly (no real computation, no array allocation, no mutation of any State value) and returns the shape/dtype structure of the output. It is built on StatefulFunction, exactly like jit(), grad(), and vmap(), so it handles existing global states the same way every other transform does.

Three behaviors are supported by a single abstract trace:

  1. Plain outputs. If f returns arrays/pytrees, the result is a pytree with jax.ShapeDtypeStruct leaves.

  2. Existing states. If f reads or writes pre-existing global State objects, they are traced transparently (no error). Their concrete values are left unchanged after the call.

  3. New nodes. If f constructs and returns a brainstate Node (e.g. lambda: brainstate.nn.LSTMCell(3, 4)), a node of the same type is reconstructed with abstract jax.ShapeDtypeStruct leaves (lazy/abstract model construction, no memory allocated). The returned node is a first-class input to subsequent brainstate transformations.

Parameters:

  • f (Callable[..., TypeVar(A)]) – The function to abstractly evaluate. It is never executed for real.

  • *args (Any) – Example positional arguments. May contain arrays, pytrees, or brainstate graph nodes.

  • return_state_shapes (bool) – If True, return (state_shapes, out_shapes) where state_shapes is a dict mapping each touched State to the jax.ShapeDtypeStruct of its value. Default is False (return only the output shapes).

  • **kwargs (Any) – Example keyword arguments.

Return type:

TypeVar(A)

Returns:

  • out_shapes (Any) – The abstract output of f (pytree of jax.ShapeDtypeStruct, or a reconstructed abstract Node). When return_state_shapes=True this is the second element of the returned (state_shapes, out_shapes) tuple.

  • state_shapes (dict) – Only when return_state_shapes=True: dict of State -> jax.ShapeDtypeStruct for every state touched by f. Returned as the first element of the (state_shapes, out_shapes) tuple.

Examples

>>> import brainstate
>>> import jax.numpy as jnp
>>> out = brainstate.transform.eval_shape(lambda x: x * 2.0, jnp.ones(3))
>>> out.shape, out.dtype
((3,), dtype('float32'))

>>> model = brainstate.transform.eval_shape(lambda: brainstate.nn.LSTMCell(3, 4))
>>> isinstance(model, brainstate.nn.LSTMCell)
True
Contents