How to count all trainable parameters for a nnx.Module

[jlperla]

I want to get a count of all of the trainable parameters for a module. In pytorch I can use model.parameters and then recursively check if a gradient is required, etc.

How would I do this with a recursive function and filters in nnx?

[cgarciae]

Hey @jlperla, the easiest way is to use nnx.state to extract the nnx.Param instances:

params = nnx.state(model, nnx.Param)
total_params = sum(np.prod(x.shape) for x in jax.tree.leaves(params), 0)

[Vilin97]

The above does not work directly for me, but this does:

params = nnx.state(model, nnx.Param)
total_params  = sum(np.prod(x.shape) for x in jax.tree_util.tree_leaves(params))