Finer granularity of control over triangular multiplicative update #205
Description
Hi,
It is a great pleasure to use cuequivariance_torch.triangle_multiplicative_update, but the all-encompassing API is a bit problmatic for us. For instance, it is easy to just set layernorm's weight and bias to 1s and 0s respectively, but not that easy if we would just want to do a rms norm instead of layernorm.
Could you guys consider providing finer control?
Best,
Rui