| 23 | |
| 24 | template <typename T, typename Context> |
| 25 | void RnnKernel(const Context& dev_ctx, |
| 26 | const DenseTensor& x, |
| 27 | const std::vector<const DenseTensor*>& pre_state, |
| 28 | const std::vector<const DenseTensor*>& weight_list, |
| 29 | const optional<DenseTensor>& sequence_length, |
| 30 | float dropout_prob, |
| 31 | bool is_bidirec, |
| 32 | int input_size, |
| 33 | int hidden_size, |
| 34 | int num_layers, |
| 35 | const std::string& mode, |
| 36 | int seed, |
| 37 | bool is_test, |
| 38 | DenseTensor* out, |
| 39 | DenseTensor* dropout_state, |
| 40 | std::vector<DenseTensor*> state, |
| 41 | DenseTensor* reserve) { |
| 42 | if (dropout_state->IsInitialized()) { |
| 43 | if (dropout_state->numel() != out->numel()) dropout_state->clear(); |
| 44 | } |
| 45 | |
| 46 | dropout_state->Resize(out->dims()); |
| 47 | dev_ctx.template Alloc<uint8_t>(dropout_state); |
| 48 | |
| 49 | funcs::SetConstant<XPUContext, uint8_t> ones; |
| 50 | ones(dev_ctx, dropout_state, static_cast<uint8_t>(1)); |
| 51 | |
| 52 | PADDLE_ENFORCE_EQ( |
| 53 | mode, |
| 54 | "LSTM", |
| 55 | errors::InvalidArgument( |
| 56 | "XPU only support LSTM mode now, current mode is %s", mode)); |
| 57 | |
| 58 | auto init_h = pre_state[0]; |
| 59 | auto init_c = pre_state[1]; |
| 60 | auto last_h = state[0]; |
| 61 | auto last_c = state[1]; |
| 62 | |
| 63 | // check shape |
| 64 | const int64_t seq_len = x.dims()[0]; // time_step |
| 65 | const int64_t batch_size = x.dims()[1]; |
| 66 | const int64_t input_dim = x.dims()[2]; |
| 67 | const int64_t direction_num = is_bidirec ? 2 : 1; |
| 68 | |
| 69 | PADDLE_ENFORCE_EQ( |
| 70 | init_h->dims()[0], |
| 71 | num_layers * direction_num, |
| 72 | errors::InvalidArgument("The num_layers of in RNN layer must" |
| 73 | " be the same as first dim of init " |
| 74 | "hidden, but received num_layers:%d," |
| 75 | " dim:%d", |
| 76 | num_layers, |
| 77 | init_h->dims()[0])); |
| 78 | |
| 79 | PADDLE_ENFORCE_EQ( |
| 80 | init_c->dims()[0], |
| 81 | num_layers * direction_num, |
| 82 | errors::InvalidArgument( |
nothing calls this directly
no test coverage detected