(self, state, embed)
| 246 | self.cast = P.Cast() |
| 247 | |
| 248 | def construct(self, state, embed): |
| 249 | state = P.Reshape()(state, (-1, self.hidden_size)) |
| 250 | # output logits over vocabulary [bs*seq_length, vocab_size] |
| 251 | logits = self.matmul( |
| 252 | self.cast(state, self.dtype), self.cast(embed, self.dtype) |
| 253 | ) |
| 254 | return logits |
| 255 | |
| 256 | |
| 257 | def set_parallel_configure_for_layer( |
nothing calls this directly
no outgoing calls
no test coverage detected