(
self,
dim: int,
dim_out: Optional[int] = None,
mult: int = 4,
glu: bool = False,
glu_mult_bias: bool = False,
swish: bool = False,
post_act_ln: bool = False,
dropout: float = 0.0,
no_bias: bool = False,
zero_init_output: bool = False,
*args,
**kwargs
)
| 75 | """ |
| 76 | |
| 77 | def __init__( |
| 78 | self, |
| 79 | dim: int, |
| 80 | dim_out: Optional[int] = None, |
| 81 | mult: int = 4, |
| 82 | glu: bool = False, |
| 83 | glu_mult_bias: bool = False, |
| 84 | swish: bool = False, |
| 85 | post_act_ln: bool = False, |
| 86 | dropout: float = 0.0, |
| 87 | no_bias: bool = False, |
| 88 | zero_init_output: bool = False, |
| 89 | *args, |
| 90 | **kwargs |
| 91 | ): |
| 92 | super().__init__() |
| 93 | inner_dim = int(dim * mult) |
| 94 | dim_out = default(dim_out, dim) |
| 95 | |
| 96 | if swish: |
| 97 | activation = nn.SiLU() |
| 98 | else: |
| 99 | activation = nn.GELU() |
| 100 | |
| 101 | if glu: |
| 102 | project_in = GLU(dim, inner_dim, activation, mult_bias=glu_mult_bias) |
| 103 | else: |
| 104 | project_in = nn.Sequential( |
| 105 | BitLinear(dim, inner_dim, bias=not no_bias, *args, **kwargs), activation |
| 106 | ) |
| 107 | if post_act_ln: |
| 108 | self.ff = nn.Sequential( |
| 109 | project_in, |
| 110 | nn.LayerNorm(inner_dim), |
| 111 | nn.Dropout(dropout), |
| 112 | BitLinear(inner_dim, dim_out, bias=not no_bias, *args, **kwargs), |
| 113 | ) |
| 114 | else: |
| 115 | self.ff = nn.Sequential( |
| 116 | project_in, |
| 117 | nn.Dropout(dropout), |
| 118 | BitLinear(inner_dim, dim_out, bias=not no_bias, *args, **kwargs), |
| 119 | ) |
| 120 | |
| 121 | # init last linear layer to 0 |
| 122 | if zero_init_output: |
| 123 | init_zero_(self.ff[-1]) |
| 124 | |
| 125 | def forward(self, x): |
| 126 | """ |
nothing calls this directly
no test coverage detected