Method
__init__
(self,
in_features,
hidden_size,
out_features=None,
act_fn="gelu_tanh",
mapping=None,
dtype=None)
Source from the content-addressed store, hash-verified
| 572 | """ |
| 573 | |
| 574 | def __init__(self, |
| 575 | in_features, |
| 576 | hidden_size, |
| 577 | out_features=None, |
| 578 | act_fn="gelu_tanh", |
| 579 | mapping=None, |
| 580 | dtype=None): |
| 581 | super().__init__() |
| 582 | if out_features is None: |
| 583 | out_features = hidden_size |
| 584 | tp_group = mapping.tp_group |
| 585 | tp_size = mapping.tp_size |
| 586 | self.linear_1 = ColumnLinear(in_features=in_features, |
| 587 | out_features=hidden_size, |
| 588 | bias=True, |
| 589 | tp_group=tp_group, |
| 590 | tp_size=tp_size, |
| 591 | dtype=dtype, |
| 592 | gather_output=False) |
| 593 | self.act_1 = ACT2FN[act_fn] |
| 594 | self.linear_2 = RowLinear(in_features=hidden_size, |
| 595 | out_features=out_features, |
| 596 | bias=True, |
| 597 | tp_group=tp_group, |
| 598 | tp_size=tp_size, |
| 599 | dtype=dtype) |
| 600 | |
| 601 | def forward(self, caption): |
| 602 | hidden_states = self.linear_1(caption) |
Callers
nothing calls this directly
Tested by
no test coverage detected