Factory function that instantiates the appropriate
SSMax* module (or NULL) from a string identifier.
Usage
create_ssmax_layer(ssmax_type, num_heads, embed_dim)
Arguments
- ssmax_type
character(1) One of:
"none", "ssmax", "ssmax-mlp",
"ssmax-mlp-elementwise", "qassmax-mlp",
"qassmax-mlp-elementwise".
- num_heads
integer(1) Number of attention heads.
- embed_dim
integer(1) Total embedding dimension
(num_heads * head_dim).
Value
An nn_module instance, or NULL when
ssmax_type = "none".