3f6c8c6c0858501abfa132caf621dba99ad7d965,se3cnn/kernel.py,SE3Kernel,combination,#SE3Kernel#Any#,356

Before Change

            begin_j = 0
            for j, mj in enumerate(self.multiplicities_in):
                si = slice(begin_i, begin_i + mi * self.dims_out[i])
                sj = slice(begin_j, begin_j + mj * self.dims_in[j])

                kij = getattr(self, "kernel_{}_{}".format(i, j))  // [beta, i, j, x, y, z]
                if kij is not None:
                    b_el = kij.size(0)

                    w = weight[weight_index: weight_index + mi * mj * b_el].view(mi, mj, b_el)  // [u, v, beta]
                    weight_index += mi * mj * b_el

                    ker = torch.einsum("uvb,bijxyz->uivjxyz", (w, kij)).contiguous()  // [u, i, v, j, x, y, z]
                    kernel[si, sj] = ker.view_as(kernel[si, sj])
                else:
                    kernel[si, sj] = 0

                begin_j += mj * self.dims_in[j]

After Change

                    ker = torch.einsum("uvb,bijxyz->uivjxyz", (w, kij)).contiguous()  // [u, i, v, j, x, y, z]
                else:
                    ker = torch.tensor(0.)

                ker = ker.view(i_diff, j_diff, self.size, self.size, self.size)

                sj_kernels.append(ker)
                begin_j += mj * self.dims_in[j]

            si_kernels.append(torch.cat(sj_kernels, dim=1))
            begin_i += mi * self.dims_out[i]

Instances