Vectorize filter bank computation for relaxed separable group convolutions

[dgcnz]

Description

Currently we generate each of the num_filter_banks weights by a loop, and we run sample_H for each of them.

gconv/gconv/gnn/kernels/kernel.py

Lines 487 to 508 in 220c2a4

	# TODO: vectorize this for loop
	weight_H = torch.stack(
	[
	self.sample_H(
	H_product_H.flatten(0, 1),
	self.weight_H[i].flatten(1, -1),
	self.grid_H,
	**self.sample_H_kwargs,
	)
	.view(
	num_in_H,
	num_out_H,
	self.in_channels // self.groups,
	self.out_channels,
	*self.weight_dims,
	)
	.transpose(0, 3)
	.transpose(1, 3)
	for i in range(self.num_filter_banks)
	],
	dim=0,
	)

This probably can be vectorized, like in the relaxed lifting example:

gconv/gconv/gnn/kernels/kernel.py

Lines 267 to 277 in 220c2a4

	weight = self.sample_Rn(
	self.weight.flatten(0, 1).repeat_interleave(H.shape[0], dim=0),
	H_product.repeat(self.num_filter_banks * self.out_channels, *product_dims),
	**self.sample_Rn_kwargs,
	).view(
	self.num_filter_banks,
	self.out_channels,
	num_H,
	self.in_channels // self.groups,
	*self.kernel_size,
	)

Tasks

• Implement a vectorized equivalent of the sequential operation
• Add test to ensure equivalence between sequential and vectorized implementations akin to https://github.com/dgcnz/gconv/blob/main/tests/test_rlkernel.py