WithTag

WithTag#

class brainstate.util.WithTag(tag)[source]#

Filter objects that have a specific tag attribute.

This filter checks if an object has a ‘tag’ attribute that matches the specified tag value. It’s commonly used to filter parameters or modules in neural networks based on their assigned tags.

Parameters:: tag (str) – The tag value to match against.

tag#

The tag value to match against.

Type:: str

Examples

>>> from brainstate.util.filter import WithTag
>>> import brainstate as bs
>>>
>>> # Create a filter for 'trainable' tag
>>> filter_trainable = WithTag('trainable')
>>>
>>> # Test with an object that has the tag
>>> class Param:
...     def __init__(self, tag):
...         self.tag = tag
>>>
>>> param1 = Param('trainable')
>>> param2 = Param('frozen')
>>>
>>> filter_trainable([], param1)
True
>>> filter_trainable([], param2)
False
>>>
>>> # Use with neural network modules
>>> class MyModule(bs.Module):
...     def __init__(self):
...         super().__init__()
...         self.weight = bs.State(bs.random.randn(10, 10))
...         self.weight.tag = 'trainable'
...         self.bias = bs.State(bs.zeros(10))
...         self.bias.tag = 'frozen'

WithTag

Contents

WithTag#