OfType

Contents

OfType#

class brainstate.util.OfType(type)[source]#

Filter objects based on their type.

This filter checks if an object is an instance of a specific type or if it has a ‘type’ attribute that is a subclass of the specified type. It’s useful for filtering specific kinds of objects in a nested structure.

Parameters:

type (type) – The type to match against.

type#

The type to match against.

Type:

type

Examples

>>> from brainstate.util.filter import OfType
>>> import numpy as np
>>> import jax.numpy as jnp
>>>
>>> # Create a filter for numpy arrays
>>> array_filter = OfType(np.ndarray)
>>>
>>> # Test with different objects
>>> array_filter([], np.array([1, 2, 3]))
True
>>> array_filter([], [1, 2, 3])
False
>>>
>>> # Filter for specific module types
>>> import brainstate as bs
>>> linear_filter = OfType(bs.nn.Linear)
>>>
>>> # Use in model filtering
>>> class Model(bs.nn.Module):
...     def __init__(self):
...         super().__init__()
...         self.linear1 = bs.nn.Linear(10, 20)
...         self.linear2 = bs.nn.Linear(20, 10)
...         self.activation = bs.nn.ReLU()
>>>
>>> # Filter all Linear layers
>>> model = Model()
>>> # linear_filter will match linear1 and linear2, not activation

See also

WithTag

Filter based on tag attributes

PathContains

Filter based on path contents

to_predicate

Convert various inputs to predicates

Notes

This filter also checks for objects that have a ‘type’ attribute, which is useful for wrapped or proxy objects that maintain type information differently.

Contents