| 23 | namespace layers { |
| 24 | |
| 25 | void BERTEmbedding::operator()(const core::Tensor &input_ids, |
| 26 | const core::Tensor &position_ids, |
| 27 | const core::Tensor &token_type_ids, |
| 28 | core::Tensor *output_tensor) const { |
| 29 | if (loguru::current_verbosity_cutoff() >= 3) { |
| 30 | std::ostringstream os; |
| 31 | os << ">>>>>>>>>>>> input_ids <<<<<<<<<<<<" << std::endl; |
| 32 | input_ids.Print<int64_t>(os); |
| 33 | os << ">>>>>>>>>>>> position_ids <<<<<<<<<<<<" << std::endl; |
| 34 | position_ids.Print<int64_t>(os); |
| 35 | os << ">>>>>>>>>>>> token_type_ids <<<<<<<<<<<<" << std::endl; |
| 36 | token_type_ids.Print<int64_t>(os); |
| 37 | LOG_S(3) << os.str(); |
| 38 | } |
| 39 | |
| 40 | TT_ENFORCE_EQ( |
| 41 | input_ids.n_dim(), 2, |
| 42 | "The input ids should be a matrix with shape [BatchSize, SeqLen]."); |
| 43 | auto batch_size = input_ids.shape(0); |
| 44 | auto seq_length = input_ids.shape(1); |
| 45 | // TODO 1. switch DeviceType::CPU 2. how should I set stride? |
| 46 | auto hidden_size = word_embedings_.shape(1); |
| 47 | |
| 48 | TT_ENFORCE(output_tensor, "The output tensor should not be nullptr."); |
| 49 | output_tensor->Reshape<float>({batch_size, seq_length, hidden_size}, |
| 50 | input_ids.device_type(), input_ids.device_id()); |
| 51 | LOG_S(3) << "Look up word embedding"; |
| 52 | |
| 53 | kernels::LookupEmbedding</*Add=*/false>(output_tensor, word_embedings_, |
| 54 | input_ids); |
| 55 | |
| 56 | LOG_S(3) << "Look up token type embedding"; |
| 57 | kernels::LookupEmbedding</*Add=*/true>(output_tensor, token_type_embeddings_, |
| 58 | token_type_ids); |
| 59 | LOG_S(3) << "Look up token position embedding"; |
| 60 | kernels::LookupEmbedding</*Add=*/true>(output_tensor, position_embeddings_, |
| 61 | position_ids); |
| 62 | |
| 63 | kernels::LayerNorm<float>(layer_norm_weights_, layer_norm_bias_, |
| 64 | output_tensor); |
| 65 | } |
| 66 | void BERTEmbedding::EnforceShapeAndType() const { |
| 67 | if (loguru::current_verbosity_cutoff() >= 3) { |
| 68 | std::ostringstream os; |
nothing calls this directly
no test coverage detected