(self, input)
| 203 | ] |
| 204 | |
| 205 | def forward(self, input): |
| 206 | |
| 207 | result = torch.empty((*input.shape[:-1], self.out_dim), dtype=input.dtype, device=input.device) |
| 208 | x, y, z = input.unbind(-1) |
| 209 | |
| 210 | result[..., 0] = self.C0 |
| 211 | if self.degree > 1: |
| 212 | result[..., 1] = -self.C1 * y |
| 213 | result[..., 2] = self.C1 * z |
| 214 | result[..., 3] = -self.C1 * x |
| 215 | if self.degree > 2: |
| 216 | xx, yy, zz = x * x, y * y, z * z |
| 217 | xy, yz, xz = x * y, y * z, x * z |
| 218 | result[..., 4] = self.C2[0] * xy |
| 219 | result[..., 5] = self.C2[1] * yz |
| 220 | result[..., 6] = self.C2[2] * (2.0 * zz - xx - yy) |
| 221 | result[..., 7] = self.C2[3] * xz |
| 222 | result[..., 8] = self.C2[4] * (xx - yy) |
| 223 | if self.degree > 3: |
| 224 | result[..., 9] = self.C3[0] * y * (3 * xx - yy) |
| 225 | result[..., 10] = self.C3[1] * xy * z |
| 226 | result[..., 11] = self.C3[2] * y * (4 * zz - xx - yy) |
| 227 | result[..., 12] = self.C3[3] * z * (2 * zz - 3 * xx - 3 * yy) |
| 228 | result[..., 13] = self.C3[4] * x * (4 * zz - xx - yy) |
| 229 | result[..., 14] = self.C3[5] * z * (xx - yy) |
| 230 | result[..., 15] = self.C3[6] * x * (xx - 3 * yy) |
| 231 | if self.degree > 4: |
| 232 | result[..., 16] = self.C4[0] * xy * (xx - yy) |
| 233 | result[..., 17] = self.C4[1] * yz * (3 * xx - yy) |
| 234 | result[..., 18] = self.C4[2] * xy * (7 * zz - 1) |
| 235 | result[..., 19] = self.C4[3] * yz * (7 * zz - 3) |
| 236 | result[..., 20] = self.C4[4] * (zz * (35 * zz - 30) + 3) |
| 237 | result[..., 21] = self.C4[5] * xz * (7 * zz - 3) |
| 238 | result[..., 22] = self.C4[6] * (xx - yy) * (7 * zz - 1) |
| 239 | result[..., 23] = self.C4[7] * xz * (xx - 3 * yy) |
| 240 | result[..., 24] = self.C4[8] * (xx * (xx - 3 * yy) - yy * (3 * xx - yy)) |
| 241 | return result |
nothing calls this directly
no outgoing calls
no test coverage detected