| 33 | super(SparseLayerNorm, self).__init__(normalized_shape, eps, elementwise_affine) |
| 34 | |
| 35 | def forward(self, input: SparseTensor) -> SparseTensor: |
| 36 | nfeats = torch.zeros_like(input.feats) |
| 37 | for k in range(input.shape[0]): |
| 38 | bfeats = input.feats[input.layout[k]] |
| 39 | bfeats = bfeats.permute(1, 0).reshape(1, input.shape[1], -1) |
| 40 | bfeats = super().forward(bfeats) |
| 41 | bfeats = bfeats.reshape(input.shape[1], -1).permute(1, 0) |
| 42 | nfeats[input.layout[k]] = bfeats |
| 43 | return input.replace(nfeats) |
| 44 | |
| 45 | |
| 46 | class SparseGroupNorm32(SparseGroupNorm): |