shared weight tpconv

[nullplay]

Hi,

I found a bug (probably?) where the tpconv with shared_weight=True does not generate correct code.
It seems the tpconv kernel generator might be missing an if-else statement to handle the shared-weight case correctly.

batchtp:

OpenEquivariance/openequivariance/templates/loop_unroll_batch.cuh

Lines 37 to 41 in f5f3d0d

	{%- if not tpp.shared_weights %}
	WEIGHT_T* w = weights + i * {{tpp.weight_numel}};
	{%- else %}
	WEIGHT_T* w = weights;
	{%- endif %}

conv:

OpenEquivariance/openequivariance/templates/loop_unroll_conv_atomic.cuh

Lines 62 to 67 in f5f3d0d

	IRREP_T* l1 = L1_in + col * {{forward_schedule.L1.dim}} + lane_id;
	IRREP_T* l2 = L2_in + i * {{forward_schedule.L2.dim}} + lane_id;
	IRREP_T* l3 = L3_out + row * {{forward_schedule.L3.dim}} + lane_id;
	WEIGHT_T* w = weights + i * {{tpp.weight_numel}};

[vbharadwaj-bk]

Good catch, thank you! We didn't handle this case because most of the models we were looking at (e.g. DiffDock, Nequip, MACE) did not use a convolution + shared weights together. We will document this.

Adding support should not be too complicated for the atomic version (we can just add the branches you pointed out for the forward and backward passes). It will be harder for the deterministic version, since all warps must deterministically reduce the gradient of the shared weight in the backward pass.

We'd love to hear if you have a use case in mind for convolution + shared weights.

[vbharadwaj-bk]

Closed with #101 , and tests included in conv_test.py!