fix: handle torch.distributions.Distribution outputs in summary() (#329)

[Mikyx-1]

Bug

torchinfo.summary() raises a TypeError when a model's forward() returns a torch.distributions.Distribution object (e.g. Categorical, Normal) instead of a plain tensor.

Root cause: LayerInfo.calculate_size() in layer_info.py handles Tensor, ndarray, dict, list, and tuple, but falls through to a hard raise TypeError for any other type — including distribution objects.

Reproduction (from #329):

import torch, torch.nn as nn, torchinfo
from torch.distributions import Categorical

class ProbabilisticModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.linear = nn.Linear(10, 5)
        self.final = nn.LogSoftmax(dim=-1)

    def forward(self, x):
        return Categorical(logits=self.final(self.linear(x)))

torchinfo.summary(ProbabilisticModel(), input_size=(1, 10))
# TypeError: Model contains a layer with an unsupported input or output type

Fix

Add an elif branch in calculate_size() for torch.distributions.Distribution:

elif isinstance(inputs, torch.distributions.Distribution):
    size = list(inputs.batch_shape + inputs.event_shape)
    elem_bytes = 0

• batch_shape + event_shape gives the meaningful output shape (e.g. [1] for Categorical with a batch of 1).
• elem_bytes = 0 since there is no single underlying tensor to measure.
• Covers all Distribution subclasses (Categorical, Normal, Bernoulli, MultivariateNormal, etc.), not just the one from the issue report.
• The existing else: raise TypeError is preserved for truly unknown types.

Tests

• test_categorical_distribution_output — verifies the exact reproduction case from Issue with torchinfo.summary() Failing on Models with torch.distributions.Categorical #329 no longer raises and reports the correct output shape [1].
• test_normal_distribution_output — verifies the fix generalises to Normal, which has a non-trivial event_shape, reporting output shape [1, 5].

Closes #329

🤖 Generated with Claude Code

[TylerYep]

Thank you for the fix!