(self, x, x_edge)
| 134 | |
| 135 | |
| 136 | def forward(self, x, x_edge): |
| 137 | |
| 138 | num_edges = len(x_edge) |
| 139 | out = [] |
| 140 | |
| 141 | # Reshape the spherical harmonics based on m (order) |
| 142 | x._m_primary(self.mappingReduced) |
| 143 | |
| 144 | # radial function |
| 145 | if self.rad_func is not None: |
| 146 | x_edge = self.rad_func(x_edge) |
| 147 | offset_rad = 0 |
| 148 | |
| 149 | # Compute m=0 coefficients separately since they only have real values (no imaginary) |
| 150 | x_0 = x.embedding.narrow(1, 0, self.mappingReduced.m_size[0]) |
| 151 | x_0 = x_0.reshape(num_edges, -1) |
| 152 | if self.rad_func is not None: |
| 153 | x_edge_0 = x_edge.narrow(1, 0, self.fc_m0.in_features) |
| 154 | x_0 = x_0 * x_edge_0 |
| 155 | x_0 = self.fc_m0(x_0) |
| 156 | |
| 157 | x_0_extra = None |
| 158 | # extract extra m0 features |
| 159 | if self.extra_m0_output_channels is not None: |
| 160 | x_0_extra = x_0.narrow(-1, 0, self.extra_m0_output_channels) |
| 161 | x_0 = x_0.narrow(-1, self.extra_m0_output_channels, (self.fc_m0.out_features - self.extra_m0_output_channels)) |
| 162 | |
| 163 | x_0 = x_0.view(num_edges, -1, self.m_output_channels) |
| 164 | #x.embedding[:, 0 : self.mappingReduced.m_size[0]] = x_0 |
| 165 | out.append(x_0) |
| 166 | offset_rad = offset_rad + self.fc_m0.in_features |
| 167 | |
| 168 | # Compute the values for the m > 0 coefficients |
| 169 | offset = self.mappingReduced.m_size[0] |
| 170 | for m in range(1, max(self.mmax_list) + 1): |
| 171 | # Get the m order coefficients |
| 172 | x_m = x.embedding.narrow(1, offset, 2 * self.mappingReduced.m_size[m]) |
| 173 | x_m = x_m.reshape(num_edges, 2, -1) |
| 174 | |
| 175 | # Perform SO(2) convolution |
| 176 | if self.rad_func is not None: |
| 177 | x_edge_m = x_edge.narrow(1, offset_rad, self.so2_m_conv[m - 1].fc.in_features) |
| 178 | x_edge_m = x_edge_m.reshape(num_edges, 1, self.so2_m_conv[m - 1].fc.in_features) |
| 179 | x_m = x_m * x_edge_m |
| 180 | x_m = self.so2_m_conv[m - 1](x_m) |
| 181 | x_m = x_m.view(num_edges, -1, self.m_output_channels) |
| 182 | #x.embedding[:, offset : offset + 2 * self.mappingReduced.m_size[m]] = x_m |
| 183 | out.append(x_m) |
| 184 | offset = offset + 2 * self.mappingReduced.m_size[m] |
| 185 | offset_rad = offset_rad + self.so2_m_conv[m - 1].fc.in_features |
| 186 | |
| 187 | out = torch.cat(out, dim=1) |
| 188 | out_embedding = SO3_Embedding( |
| 189 | 0, |
| 190 | x.lmax_list.copy(), |
| 191 | self.m_output_channels, |
| 192 | device=x.device, |
| 193 | dtype=x.dtype |
nothing calls this directly
no test coverage detected