HiData

class brainstate.nn.HiData(children=None, name=None, **kwargs)

Hierarchical state container for composed dynamics.

Stores child states in a dictionary where keys match the attribute names of child dynamics in the parent dynamics class.

Supports two initialization styles:

• Data(children={‘key1’: data1, ‘key2’: data2})

• Data(key1=data1, key2=data2)

And two access styles:

• cd[‘key1’] or cd.key1

children

Dict mapping child names to their states.

Examples

Create a simple Data object:

>>> data = HiData(name='config', learning_rate=0.01, batch_size=32)
>>> print(data)
ParamData(
  name='config',
  learning_rate=0.01,
  batch_size=32
)

Create nested Data objects:

>>> import numpy as np
>>> optimizer = HiData(name='optimizer', lr=0.001, momentum=0.9)
>>> model = HiData(name='model', weights=np.array([1, 2, 3]))
>>> config = HiData(name='config', optimizer=optimizer, model=model)
>>> print(config)
ParamData(
  name='config',
  optimizer=ParamData(
    name='optimizer',
    lr=0.001,
    momentum=0.9
  ),
  model=ParamData(
    name='model',
    weights=Array(shape=(3,), dtype=int64)
  )
)

Access children using attribute or dictionary syntax:

>>> data = HiData(name='test', value=42)
>>> data.value
42
>>> data['value']
42

Clone and modify:

>>> original = HiData(name='original', x=1, y=2)
>>> cloned = original.clone()
>>> cloned['z'] = 3

clone()

Create a deep copy of the state, recursively cloning children.

Return type:

HiData

Returns:

New state instance with cloned tensors.

property dtype

Return dtype of first array child.

classmethod from_dict(d)

Create state from dictionary.

Parameters:

d (Dict) – Dictionary mapping state variable names to tensors.

Return type:

HiData

Returns:

State instance.

items()

Return child items.

keys()

Return child keys.

replace(**updates) → T

Replace specified fields with new values.

Return type:

TypeVar(T)

property state_size: int

Number of state variables per node.

to_dict()

Convert to dictionary representation.

Return type:

Dict

Returns:

Dictionary mapping state variable names to tensors.

values()

Return child values.