| 142 | } |
| 143 | |
| 144 | TfLiteStatus EvalFloat(const TfLiteTensor* input, |
| 145 | const TfLiteTensor* input_weights, |
| 146 | const TfLiteTensor* recurrent_weights, |
| 147 | const TfLiteTensor* bias, |
| 148 | const TfLiteSequenceRNNParams* params, |
| 149 | TfLiteTensor* hidden_state, TfLiteTensor* output) { |
| 150 | // Initialize the pointer bias. |
| 151 | const float* bias_ptr = bias->data.f; |
| 152 | |
| 153 | const bool time_major = params->time_major; |
| 154 | const int batch_size = |
| 155 | (time_major) ? input->dims->data[1] : input->dims->data[0]; |
| 156 | const int max_time = |
| 157 | (time_major) ? input->dims->data[0] : input->dims->data[1]; |
| 158 | const int num_units = input_weights->dims->data[0]; |
| 159 | const int input_size = input->dims->data[2]; |
| 160 | |
| 161 | // Initialize input_weights and recurrent_weights. |
| 162 | const float* input_weights_ptr = input_weights->data.f; |
| 163 | const float* recurrent_weights_ptr = recurrent_weights->data.f; |
| 164 | |
| 165 | if (time_major) { |
| 166 | // Initialize the pointer to hidden state. |
| 167 | float* hidden_state_ptr_batch = hidden_state->data.f; |
| 168 | // Unroll the sequence and use batch operations for efficiency. |
| 169 | for (int s = 0; s < max_time; s++) { |
| 170 | // Initialize the pointer to input and output. |
| 171 | const float* input_ptr_batch = |
| 172 | input->data.f + s * input_size * batch_size; |
| 173 | float* output_ptr_batch = output->data.f + s * num_units * batch_size; |
| 174 | |
| 175 | kernel_utils::RnnBatchStep( |
| 176 | input_ptr_batch, input_weights_ptr, recurrent_weights_ptr, bias_ptr, |
| 177 | input_size, num_units, batch_size, num_units, params->activation, |
| 178 | hidden_state_ptr_batch, output_ptr_batch); |
| 179 | } |
| 180 | } else { |
| 181 | // For each batch |
| 182 | for (int b = 0; b < batch_size; b++) { |
| 183 | // Initialize the pointer to hidden state. |
| 184 | float* hidden_state_ptr_batch = hidden_state->data.f + b * num_units; |
| 185 | for (int s = 0; s < max_time; s++) { |
| 186 | // Initialize the pointer to input and output. |
| 187 | const float* input_ptr_batch = |
| 188 | input->data.f + b * input_size * max_time + s * input_size; |
| 189 | float* output_ptr_batch = |
| 190 | output->data.f + b * num_units * max_time + s * num_units; |
| 191 | |
| 192 | kernel_utils::RnnBatchStep( |
| 193 | input_ptr_batch, input_weights_ptr, recurrent_weights_ptr, bias_ptr, |
| 194 | input_size, num_units, /*batch_size=*/1, num_units, |
| 195 | params->activation, hidden_state_ptr_batch, output_ptr_batch); |
| 196 | } |
| 197 | } |
| 198 | } |
| 199 | return kTfLiteOk; |
| 200 | } |
| 201 | |