One complex parameter should count as two params

[scaomath]

As all models' parameters counting traces back here

torchinfo/torchinfo/layer_info.py

Lines 154 to 170 in 8b3ae72

	def get_param_count(
	module: nn.Module, name: str, param: torch.Tensor
	) -> tuple[int, str]:
	"""
	Get count of number of params, accounting for mask.
	Masked models save parameters with the suffix "_orig" added.
	They have a buffer ending with "_mask" which has only 0s and 1s.
	If a mask exists, the sum of 1s in mask is number of params.
	"""
	if name.endswith("_orig"):
	without_suffix = name[:-5]
	pruned_weights = rgetattr(module, f"{without_suffix}_mask")
	if pruned_weights is not None:
	parameter_count = int(torch.sum(pruned_weights))
	return parameter_count, without_suffix
	return param.nelement(), name

there is no checking on whether the parameter tensor is complex or real. If a parameter is complex, such as a + i b, then it represents actually two parameters (for counting MACs/FLOPs purpose).

Of course this PR might not be conforming with torchinfo's dev, feel free to close it, I hope complex would be considered in next version.

[TylerYep]

Thanks for looking into this. Could you provide a small example model + output to illustrate your point? That will help validate this fix works

[scaomath]

Thanks for looking into this. Could you provide a small example model + output to illustrate your point? That will help validate this fix works

I added a test. Complex layer correctly returns the double params count. I encounter this issue when running torchinfo for network like FNO: neuraloperator/neuraloperator#52
I thought it would be nice to integrate this to one of my all time fav add-on for torch.

Thoughts? @TylerYep