:param x: tuple (s, V) of `torch.Tensor`, or single `torch.Tensor` (will be assumed to be scalar channels)
(self, x)
| 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): |
nothing calls this directly
no test coverage detected