Function EvalFloat

tensorflow/lite/kernels/bidirectional_sequence_rnn.cc:271–403

}

TfLiteStatus EvalFloat(const TfLiteTensor* input, const TfLiteTensor* bw_input,
                       const TfLiteTensor* fw_input_weights,
                       const TfLiteTensor* fw_recurrent_weights,
                       const TfLiteTensor* fw_bias,
                       const TfLiteTensor* bw_input_weights,
                       const TfLiteTensor* bw_recurrent_weights,
                       const TfLiteTensor* bw_bias,
                       const TfLiteTensor* aux_input,
                       const TfLiteTensor* fw_aux_input_weights,
                       const TfLiteTensor* bw_aux_input_weights,
                       const TfLiteBidirectionalSequenceRNNParams* params,
                       TfLiteTensor* fw_hidden_state, TfLiteTensor* fw_output,
                       TfLiteTensor* bw_hidden_state, TfLiteTensor* bw_output) {
  const bool time_major = params->time_major;
  const int batch_size =
      (time_major) ? input->dims->data[1] : input->dims->data[0];
  const int max_time =
      (time_major) ? input->dims->data[0] : input->dims->data[1];
  const int input_size = input->dims->data[2];
  const int aux_input_size = (aux_input) ? aux_input->dims->data[2] : 0;

  const int fw_num_units = fw_input_weights->dims->data[0];
  const float* fw_bias_ptr = fw_bias->data.f;
  const float* fw_input_weights_ptr = fw_input_weights->data.f;
  const float* fw_recurrent_weights_ptr = fw_recurrent_weights->data.f;

  const int bw_num_units = bw_input_weights->dims->data[0];
  const float* bw_bias_ptr = bw_bias->data.f;
  const float* bw_input_weights_ptr = bw_input_weights->data.f;
  const float* bw_recurrent_weights_ptr = bw_recurrent_weights->data.f;

  const float* fw_aux_input_weights_ptr = (fw_aux_input_weights != nullptr)
                                              ? fw_aux_input_weights->data.f
                                              : nullptr;
  const float* bw_aux_input_weights_ptr = (bw_aux_input_weights != nullptr)
                                              ? bw_aux_input_weights->data.f
                                              : nullptr;

  const int fw_output_step =
      params->merge_outputs ? fw_num_units + bw_num_units : fw_num_units;
  const int bw_output_step =
      params->merge_outputs ? fw_num_units + bw_num_units : bw_num_units;
  if (time_major) {
    // Forward cell.
    float* fw_hidden_state_ptr_batch = fw_hidden_state->data.f;
    for (int s = 0; s < max_time; s++) {
      const float* input_ptr_batch =
          input->data.f + s * input_size * batch_size;
      const float* aux_input_ptr_batch =
          (aux_input != nullptr)
              ? aux_input->data.f + s * input_size * batch_size
              : nullptr;
      float* output_ptr_batch =
          fw_output->data.f + s * fw_output_step * batch_size;

      kernel_utils::RnnBatchStep(
          input_ptr_batch, fw_input_weights_ptr, aux_input_ptr_batch,
          fw_aux_input_weights_ptr, fw_recurrent_weights_ptr, fw_bias_ptr,