(g, layer, batch_averaged)
| 152 | |
| 153 | @staticmethod |
| 154 | def conv2d(g, layer, batch_averaged): |
| 155 | # g: batch_size * n_filters * out_h * out_w |
| 156 | # n_filters is actually the output dimension (analogous to Linear layer) |
| 157 | spatial_size = g.size(2) * g.size(3) |
| 158 | batch_size = g.shape[0] |
| 159 | g = g.transpose(1, 2).transpose(2, 3) |
| 160 | g = try_contiguous(g) |
| 161 | g = g.view(-1, g.size(-1)) |
| 162 | |
| 163 | if batch_averaged: |
| 164 | g = g * batch_size |
| 165 | g = g * spatial_size |
| 166 | cov_g = g.t() @ (g / g.size(0)) |
| 167 | |
| 168 | return cov_g |
| 169 | |
| 170 | @staticmethod |
| 171 | def linear(g, layer, batch_averaged): |
nothing calls this directly
no test coverage detected