brainstate.nn.count_parameters

brainstate.nn.count_parameters(module, precision=2, return_table=False)

Count and display the number of trainable parameters in a neural network model.

This function iterates through all the parameters of the given model, counts the number of parameters for each module, and displays them in a table. It also calculates and returns the total number of trainable parameters.

Parameters:

model (brainstate.nn.Module) – The neural network model for which to count parameters.

Returns:

The total number of trainable parameters in the model.

Return type:

int | Tuple[Any, int]

Notes

Prints a pretty-formatted table showing the number of parameters for each module, followed by the total number of trainable parameters.