probly.traverse_nn.utils.get_output_dim

probly.traverse_nn.utils.get_output_dim(model: object) → int

Return the output feature dimension of a model.

Walks model using the neural-network traverser and returns the output feature dimension of the last weight-bearing layer visited in forward DFS order. For typical classifier-style encoders (e.g., torchvision ResNet) this corresponds to the final Linear layer’s out_features or the final Conv layer’s out_channels.

The walk does not mutate or deep-copy the input model. Backend-specific handlers must be registered on output_dim_traverser for each layer type that should contribute an output dim (for torch, this is done in probly.utils.torch).

Parameters:

model – The model whose output dim to infer. Must be of a type for which handlers have been registered (currently torch nn.Module).

Returns:

The number of output features of the model’s last weight-bearing layer.

Raises:

ValueError – If no weight-bearing layer was visited during traversal (e.g., the model contains no Linear or Conv layer, or no handlers are registered for the model’s backend).