def message(self, x_j, pseudo):
data = SplineBasis.apply(
pseudo,
self._buffers["kernel_size"],
self._buffers["is_open_spline"],
self.degree,
)
z_j = SplineWeighting.apply(x_j, self.weight, *data)
return self.ball_out.expmap0(z_j)
def forward(self, pseudo, ref_idx, y, ref_deg, batch_size):
conv_out = y[0].new_zeros(self.in_channels, self.ref_size * batch_size,
self.out_channels)
for c in range(self.in_channels):
data = SplineBasis.apply(pseudo[c], self.kernel_size,
self.is_open_spline, self.spline_degree)
out = SplineWeighting.apply(y[c], self.weight[:, c].unsqueeze(1),
*data)
idx = ref_idx[c].expand_as(out)
conv_out[c].scatter_add_(0, idx, out)
if self.norm:
conv_out[c].div_(ref_deg[c])
conv_out = conv_out.sum(0)
conv_out = conv_out.view(batch_size, self.ref_size, self.out_channels)
conv_out = conv_out.transpose(1, 2)
if self.bias is not None:
conv_out = conv_out + self.bias[:, None]
return conv_out
def message(self, x_j, pseudo):
data = SplineBasis.apply(pseudo, self._buffers['kernel_size'],
self._buffers['is_open_spline'], self.degree)
return SplineWeighting.apply(x_j, self.weight, *data)