Order of dimensions in module vs expand

Misc.
1

In the Bayesian regression example:

class BayesianRegression(PyroModule):
    def __init__(self, in_features, out_features):
        super().__init__()
        self.linear = PyroModule[nn.Linear](in_features, out_features)
        self.linear.weight = PyroSample(dist.Normal(0., 1.).expand([out_features, in_features]).to_event(2))
        self.linear.bias = PyroSample(dist.Normal(0., 10.).expand([out_features]).to_event(1))

I am confused why in the definition of the linear layer the order is (in_features, out_features) but in the weight definition we do expand([out_features, in_features]). I’ve been trying to create my own simple NN and I’m getting all kinds of errors related to incorrect dimensions, and I can’t seem to figure it out.

2

It comes from the definition of PyTorch nn.Linear (you can scroll down a bit to see the definition of Linear.weight). There, the constructor requires in, out and the weight has shape out, in.

3

Hmm ok. That was really confusing me. Thanks!