Question about LayerNorm in VectorizedLinearLayer

[ZiyiLiubird]

You implement LayerNorm in VectorizedLinearLayer but the LayerNorm input is "self._out_features, self._population_size", which means that the eps in LayerNorm is self._population_size. Could you please explain why ?

        self._layer_norm = (
            torch.nn.LayerNorm(self._out_features, self._population_size)
            if use_layer_norm
            else None
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        assert x.shape[0] == self._population_size
        if self._layer_norm is not None:
            return self._layer_norm(x.matmul(self.weight) + self.bias)
        return x.matmul(self.weight) + self.bias