find_layers¶

probly.traverse_nn.utils.find_layers(model: object, layer_types: type | tuple[type, ...]) → list[Any][source]¶

Return all layers of the given type(s) contained in a model.

Walks model using the neural-network traverser and collects every visited layer that is an instance of layer_types, in forward DFS order. The check is a plain isinstance, so no per-backend handlers are needed; any backend supported by nn_traverser (currently torch and flax NNX) works out of the box. The walk does not mutate or deep-copy the model.

Parameters:

model – The model to search.
layer_types – A layer type or tuple of layer types to match.

Returns:

All matching layers in traversal order; empty if none match.