Method
__init__
(
self,
dim_in: int,
dim_out: int,
activation: Callable,
mult_bias: bool = False,
linear: Callable = False,
*args,
**kwargs
)
Source from the content-addressed store, hash-verified
| 28 | """ |
| 29 | |
| 30 | def __init__( |
| 31 | self, |
| 32 | dim_in: int, |
| 33 | dim_out: int, |
| 34 | activation: Callable, |
| 35 | mult_bias: bool = False, |
| 36 | linear: Callable = False, |
| 37 | *args, |
| 38 | **kwargs |
| 39 | ): |
| 40 | super().__init__() |
| 41 | self.dim_in = dim_in |
| 42 | self.dim_out = dim_out |
| 43 | self.activation = activation |
| 44 | self.mult_bias = mult_bias |
| 45 | |
| 46 | if linear: |
| 47 | self.proj = linear(dim_in, dim_out * 2) |
| 48 | else: |
| 49 | self.proj = BitLinear(dim_in, dim_out * 4, *args, **kwargs) |
| 50 | |
| 51 | self.mult_bias = nn.Parameter(torch.ones(dim_out)) if mult_bias else 1.0 |
| 52 | |
| 53 | def forward(self, x: Tensor): |
| 54 | x, gate = self.proj(x).chunk(2, dim=-1) |
Tested by
no test coverage detected