Combined LayerNorm for tuples (s, V). Takes tuples (s, V) as input and as output.
| 200 | |
| 201 | |
| 202 | class GVPLayerNorm(nn.Module): |
| 203 | ''' |
| 204 | Combined LayerNorm for tuples (s, V). |
| 205 | Takes tuples (s, V) as input and as output. |
| 206 | ''' |
| 207 | def __init__(self, dims): |
| 208 | super().__init__() |
| 209 | self.s, self.v = dims |
| 210 | self.scalar_norm = nn.LayerNorm(self.s) |
| 211 | |
| 212 | def forward(self, x): |
| 213 | ''' |
| 214 | :param x: tuple (s, V) of `torch.Tensor`, |
| 215 | or single `torch.Tensor` |
| 216 | (will be assumed to be scalar channels) |
| 217 | ''' |
| 218 | if not self.v: |
| 219 | return self.scalar_norm(x) |
| 220 | s, v = x |
| 221 | vn = _norm_no_nan(v, axis=-1, keepdims=True, sqrt=False) |
| 222 | vn = torch.sqrt(torch.mean(vn, dim=-2, keepdim=True)) |
| 223 | return self.scalar_norm(s), v / vn |
| 224 | |
| 225 | |
| 226 | class GVPConv(MessagePassing): |