Array

class brainstate.typing.Array

Flexible array type annotation supporting shape and dtype specifications.

This class provides a convenient way to annotate arrays with shape information, making code more self-documenting and enabling better static analysis.

Examples

Basic array annotations:

>>> from brainstate.typing import Array
>>>
>>> # Any array
>>> def process_array(x: Array) -> Array:
...     return x * 2
>>>
>>> # Array with specific shape annotation
>>> def matrix_multiply(a: Array["m, n"], b: Array["n, k"]) -> Array["m, k"]:
...     return a @ b
>>>
>>> # Array with dtype and shape
>>> def normalize_weights(weights: Array["batch, features"]) -> Array["batch, features"]:
...     return weights / weights.sum(axis=-1, keepdims=True)

Advanced shape annotations:

>>> # Using ellipsis for flexible dimensions
>>> def flatten_batch(x: Array["batch, ..."]) -> Array["batch, -1"]:
...     return x.reshape(x.shape[0], -1)
>>>
>>> # Multiple shape constraints
>>> def attention(
...     query: Array["batch, seq_len, d_model"],
...     key: Array["batch, seq_len, d_model"],
...     value: Array["batch, seq_len, d_model"]
... ) -> Array["batch, seq_len, d_model"]:
...     # Attention computation
...     pass