Jake, sorry for the late reply. if your input is B (an even number equal to 2xb) and your ooutput is C, and if weight sharing is acceptable, could you just reshape to 2 x b and then simply apply a linear layer (b,C) to the last dimension?