probly.traverse_nn.common

Generic traverser helpers for neural networks.

Functions

compose(*traversers[, nn_traverser, name])

Compose a custom traverser with neural network traversal functionality.
probly.traverse_nn.common.compose(*traversers, nn_traverser=<pytraverse.composition.LazydispatchTraverser object>, name=None)[source]

Compose a custom traverser with neural network traversal functionality.

This function creates a sequential traverser that combines neural network traversal, a custom traverser, and layer counting capabilities in a specific order.

Parameters:

  • traversers (Traverser) – Custom traverser functions to be composed (sequentially) with the NN traverser.

  • nn_traverser (Traverser) – The neural network traverser to use. Defaults to the module’s nn_traverser.

  • name (str | None) – Optional name for the composed traverser.

Returns:

A composed sequential traverser that applies NN traversal, custom traversal, and layer counting in sequence.

Return type:

Traverser