Key

Contents

Key#

class brainstate.typing.Key(*args, **kwargs)[source]#

Protocol for keys that can be used in PyTree paths.

A Key must be both hashable and comparable, making it suitable for use as dictionary keys and for ordering operations.

Examples

Valid key types include:

>>> # String keys
>>> key1: Key = "layer1"
>>>
>>> # Integer keys
>>> key2: Key = 42
>>>
>>> # Custom hashable objects
>>> class CustomKey:
...     def __init__(self, name: str):
...         self.name = name
...
...     def __hash__(self) -> int:
...         return hash(self.name)
...
...     def __eq__(self, other) -> bool:
...         return isinstance(other, CustomKey) and self.name == other.name
...
...     def __lt__(self, other) -> bool:
...         return isinstance(other, CustomKey) and self.name < other.name
Contents