// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.

#include "core/providers/cpu/rnn/rnn_helpers.h"

#include <cmath>
#include <cstring>
#include <iomanip>
#include <iostream>
#include <stdlib.h>
#include <string>
#include <unordered_map>

#include "core/common/common.h"
#include "core/framework/op_kernel.h"
#include "core/providers/cpu/rnn/rnn_activation_functors.h"
#include "core/util/math.h"
#include "core/util/math_cpuonly.h"

namespace onnxruntime {
namespace rnn {
namespace detail {

using namespace ::onnxruntime::common;

Status ValidateCommonRnnInputs(const Tensor& X,
                               const Tensor& W,
                               const Tensor& R,
                               const Tensor* B,
                               int WRB_dim_1_multipler,
                               const Tensor* sequence_lens,
                               const Tensor* initial_h,
                               int64_t num_directions,
                               int64_t hidden_size) {
  auto& X_shape = X.Shape();
  auto& W_shape = W.Shape();
  auto& R_shape = R.Shape();

  int64_t seq_length = X_shape[0];
  int64_t batch_size = X_shape[1];
  int64_t input_size = X_shape[2];

  if (X_shape.NumDimensions() != 3)
    return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Input X must have 3 dimensions only. Actual:", X_shape);

  if (W_shape.NumDimensions() != 3 ||
      W_shape[0] != num_directions ||
      W_shape[1] != hidden_size * WRB_dim_1_multipler ||
      W_shape[2] != input_size)
    return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Input W must have shape {",
                           num_directions, ",", WRB_dim_1_multipler, "*", hidden_size, ",",
                           input_size, "}. Actual:", W_shape);

  if (R_shape.NumDimensions() != 3 ||
      R_shape[0] != num_directions ||
      R_shape[1] != hidden_size * WRB_dim_1_multipler ||
      R_shape[2] != hidden_size)
    return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Input R must have shape {",
                           num_directions, ",", WRB_dim_1_multipler, "*", hidden_size, ",",
                           hidden_size, "}. Actual:", R_shape);

  if (B != nullptr) {
    auto& B_shape = B->Shape();
    if (B_shape.NumDimensions() != 2 ||
        B_shape[0] != num_directions ||
        B_shape[1] != 2 * WRB_dim_1_multipler * hidden_size)
      return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Input B must have shape {",
                             num_directions, ",", 2 * WRB_dim_1_multipler, "*", hidden_size, "}. Actual:", B_shape);
  }

  if (sequence_lens != nullptr) {
    auto& sequence_lens_shape = sequence_lens->Shape();
    if (sequence_lens_shape.NumDimensions() != 1 ||
        sequence_lens_shape[0] != batch_size) {
      return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Input sequence_lens must have shape {",
                             batch_size, "}. Actual:", sequence_lens_shape);
    }

    auto sequence_len_entries = sequence_lens->DataAsSpan<int>();
    if (std::any_of(sequence_len_entries.cbegin(),
                    sequence_len_entries.cend(),
                    [seq_length](int len) { return len <= 0 || len > seq_length; })) {
      return ORT_MAKE_STATUS(
          ONNXRUNTIME, INVALID_ARGUMENT,
          "Invalid value/s in sequence_lens. All values must be > 0 and < seq_length. seq_length=", seq_length);
    }
  }

  if (initial_h != nullptr) {
    auto& initial_h_shape = initial_h->Shape();

    if (initial_h_shape.NumDimensions() != 3 ||
        initial_h_shape[0] != num_directions ||
        initial_h_shape[1] != batch_size ||
        initial_h_shape[2] != hidden_size)

      return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "Input initial_h must have shape {",
                             num_directions, ",", batch_size, ",", hidden_size, "}. Actual:", initial_h_shape);
  }

  return Status::OK();
}  // namespace detail

// map of arg name and whether the alpha and/or beta arguments are required
static std::unordered_map<std::string, std::pair<bool, bool>>
    NameToArgUsageMap{{"affine", {1, 1}},
                      {"relu", {0, 0}},
                      {"leakyrelu", {1, 0}},
                      {"thresholdedrelu", {1, 0}},
                      {"tanh", {0, 0}},
                      {"scaledtanh", {1, 1}},
                      {"sigmoid", {0, 0}},
                      {"hardsigmoid", {1, 1}},
                      {"elu", {1, 0}},
                      {"softsign", {0, 0}},
                      {"softplus", {0, 0}}};

// map of alpha/beta defaults
static std::unordered_map<std::string, std::pair<float, float>>
    NameToArgDefaultsMap{{"leakyrelu", {0.01f, 0.f}},
                         {"hardsigmoid", {0.2f, 0.5f}},
                         {"elu", {1.0f, 0.f}}};

std::string NormalizeActivationArgumentAndGetAlphaBetaCount(const std::string& activation,
                                                            std::vector<float>::const_iterator& cur_alpha,
                                                            const std::vector<float>::const_iterator& end_alpha,
                                                            std::vector<float>::const_iterator& cur_beta,
                                                            const std::vector<float>::const_iterator& end_beta,
                                                            float& alpha, float& beta) {
  std::string name(activation);
  std::transform(name.begin(), name.end(), name.begin(),
                 [](const unsigned char i) { return static_cast<char>(::tolower(i)); });

  auto usage_entry = NameToArgUsageMap.find(name);
  if (usage_entry == NameToArgUsageMap.end()) {
    ORT_THROW(
        "Expecting activation to be one of Affine, Relu, LeakyRelu, "
        "ThresholdedRelu, Tanh, ScaledTanh, Sigmoid, HardSigmoid, "
        "Elu, Softsign, Softplus. Got " +
        activation);
  }

  bool needs_alpha = usage_entry->second.first;
  bool needs_beta = usage_entry->second.second;

  auto set_if_needed = [](bool needed,
                          std::vector<float>::const_iterator& in,
                          const std::vector<float>::const_iterator& in_end,
                          const float default_val,
                          float& out) {
    if (needed) {
      if (in != in_end) {
        out = *in;
        ++in;
      } else {
        out = default_val;
      }
    }
  };

  auto defaults_entry = NameToArgDefaultsMap.find(name);
  if (defaults_entry != NameToArgDefaultsMap.end()) {
    set_if_needed(needs_alpha, cur_alpha, end_alpha, defaults_entry->second.first, alpha);
    set_if_needed(needs_beta, cur_beta, end_beta, defaults_entry->second.second, beta);
  } else {
    set_if_needed(needs_alpha, cur_alpha, end_alpha, 0.f, alpha);
    set_if_needed(needs_beta, cur_beta, end_beta, 0.f, beta);
  }

  return name;
}

ActivationFuncs::ActivationFuncs(const std::vector<std::string>& funcs,
                                 const std::vector<float>& alphas,
                                 const std::vector<float>& betas) {
  auto cur_alpha = alphas.cbegin();
  auto end_alpha = alphas.cend();
  auto cur_beta = betas.cbegin();
  auto end_beta = betas.cend();

  for (auto input_func : funcs) {
    float alpha = 0.f;
    float beta = 0.f;
    std::string func = detail::NormalizeActivationArgumentAndGetAlphaBetaCount(
        input_func, cur_alpha, end_alpha, cur_beta, end_beta, alpha, beta);

    entries_.push_back(Entry{func, alpha, beta});
  }
}

void DumpMatrixImpl(const std::string& name, const float* src, int row, int col, int offset, int col_width) {
  std::cout << "Dump matrix: " << name << std::endl;

  if (col_width == -1) col_width = col;

  for (int r = 0; r < row; r++) {
    for (int c = 0; c < col; c++) {
      int index = r * col_width + offset + c;
      std::cout << std::setw(12) << std::setprecision(8) << src[index];
    }
    std::cout << std::endl;
  }
  std::cout << std::endl;
}

namespace deepcpu {

const float alpha_1 = 4.89352455891786e-03f;
const float alpha_3 = 6.37261928875436e-04f;
const float alpha_5 = 1.48572235717979e-05f;
const float alpha_7 = 5.12229709037114e-08f;
const float alpha_9 = -8.60467152213735e-11f;
const float alpha_11 = 2.00018790482477e-13f;
const float alpha_13 = -2.76076847742355e-16f;

const float beta_0 = 4.89352518554385e-03f;
const float beta_2 = 2.26843463243900e-03f;
const float beta_4 = 1.18534705686654e-04f;
const float beta_6 = 1.19825839466702e-06f;

const float sigmoid_bound = 20.0f;
const float tanh_bound = 10.0f;

inline void clip_for_sigmoid_in_place(float* ps, int c) {
  for (int i = 0; i < c; i++) {
    if (ps[i] < -sigmoid_bound)
      ps[i] = -sigmoid_bound;
    else if (ps[i] > sigmoid_bound)
      ps[i] = sigmoid_bound;
  }
}

inline void clip_for_tanh_in_place(float* ps, int c) {
  for (int i = 0; i < c; i++) {
    if (ps[i] < -tanh_bound)
      ps[i] = -tanh_bound;
    else if (ps[i] > tanh_bound)
      ps[i] = tanh_bound;
  }
}

inline void clip_for_sigmoid(const float* ps, float* pd, int c) {
  for (int i = 0; i < c; i++) {
    if (ps[i] < -sigmoid_bound)
      pd[i] = -sigmoid_bound;
    else if (ps[i] > sigmoid_bound)
      pd[i] = sigmoid_bound;
    else
      pd[i] = ps[i];
  }
}

inline void clip_for_tanh(const float* ps, float* pd, int c) {
  for (int i = 0; i < c; i++) {
    if (ps[i] < -tanh_bound)
      pd[i] = -tanh_bound;
    else if (ps[i] > tanh_bound)
      pd[i] = tanh_bound;
    else
      pd[i] = ps[i];
  }
}

void add_bias_into_ignore(const float* ps, float* pd, const int c) {
  ORT_UNUSED_PARAMETER(ps);
  ORT_UNUSED_PARAMETER(pd);
  ORT_UNUSED_PARAMETER(c);
}

void add_bias_into(const float* ps, float* pd, const int c) {
  for (int i = 0; i < c; i++) {
    pd[i] += ps[i];
  }
}

void clip(const float b, float* pd, const int c) {
  for (int i = 0; i < c; i++) {
    float x = pd[i];
    if (x > b)
      pd[i] = b;
    else if (x < -b)
      pd[i] = -b;
  }
}

void clip_ignore_bias(const float b, const float* pb, float* pd, const int c) {
  ORT_UNUSED_PARAMETER(pb);

  for (int i = 0; i < c; i++) {
    float x = pd[i];
    if (x > b)
      pd[i] = b;
    else if (x < -b)
      pd[i] = -b;
    else
      pd[i] = x;
  }
}

void clip_add_bias(const float b, const float* pb, float* pd, const int c) {
  for (int i = 0; i < c; i++) {
    float x = pd[i] + pb[i];
    if (x > b)
      pd[i] = b;
    else if (x < -b)
      pd[i] = -b;
    else
      pd[i] = x;
  }
}

void sigmoid_m(const float* ps1, float* ps1_c, const float* ps2, float* pd, int c,
               const float alpha, const float beta) {
  ORT_UNUSED_PARAMETER(alpha);
  ORT_UNUSED_PARAMETER(beta);

  clip_for_sigmoid(ps1, ps1_c, c);

  for (int i = 0; i < c; i++) {
    float x = 0.5f * ps1_c[i];
    float x2 = x * x;
    float p = x2 * alpha_13 + alpha_11;
    p = x2 * p + alpha_9;
    p = x2 * p + alpha_7;
    p = x2 * p + alpha_5;
    p = x2 * p + alpha_3;
    p = x2 * p + alpha_1;
    p = x * p;
    float q = x2 * beta_6 + beta_4;
    q = x2 * q + beta_2;
    q = x2 * q + beta_0;
    pd[i] = ps2[i] * 0.5f * (1 + (p / q));
  }
}

void tanh_m(const float* ps1, float* ps1_c, const float* ps2, float* pd, int c,
            const float alpha, const float beta) {
  ORT_UNUSED_PARAMETER(alpha);
  ORT_UNUSED_PARAMETER(beta);

  clip_for_tanh(ps1, ps1_c, c);

  for (int i = 0; i < c; i++) {
    float x = ps1_c[i];
    float x2 = x * x;
    float p = x2 * alpha_13 + alpha_11;
    p = x2 * p + alpha_9;
    p = x2 * p + alpha_7;
    p = x2 * p + alpha_5;
    p = x2 * p + alpha_3;
    p = x2 * p + alpha_1;
    p = x * p;
    float q = x2 * beta_6 + beta_4;
    q = x2 * q + beta_2;
    q = x2 * q + beta_0;
    pd[i] = ps2[i] * p / q;
  }
}

void relu_m(const float* ps1, float* ps1_c, const float* ps2, float* pd, int c,
            const float alpha, const float beta) {
  ORT_UNUSED_PARAMETER(ps1_c);
  ORT_UNUSED_PARAMETER(alpha);
  ORT_UNUSED_PARAMETER(beta);

  for (int i = 0; i < c; i++) {
    const float max = ps1[i] > 0 ? ps1[i] : 0.0f;
    pd[i] = ps2[i] * max;
  }
}

void composed_m(const float* ps1, float* ps1_c, const float* ps2, float* pd, int c,
                std::function<float(float, float, float)> func,
                const float alpha, const float beta) {
  ORT_UNUSED_PARAMETER(ps1_c);
  for (int i = 0; i < c; i++) {
    pd[i] = ps2[i] * func(ps1[i], alpha, beta);
  }
}

void sigmoid_exact_m(const float* ps1, float* ps1_c, const float* ps2, float* pd, int c,
                     const float alpha, const float beta) {
  ORT_UNUSED_PARAMETER(ps1_c);
  ORT_UNUSED_PARAMETER(alpha);
  ORT_UNUSED_PARAMETER(beta);

  for (int i = 0; i < c; i++) {
    float x = ps1[i];
    pd[i] = ps2[i] / (1 + ::std::exp(-x));
  }
}

void tanh_exact_m(const float* ps1, float* ps1_c, const float* ps2, float* pd, int c,
                  const float alpha, const float beta) {
  ORT_UNUSED_PARAMETER(ps1_c);
  ORT_UNUSED_PARAMETER(alpha);
  ORT_UNUSED_PARAMETER(beta);

  for (int i = 0; i < c; i++) {
    pd[i] = ::std::tanh(ps1[i]) * ps2[i];
  }
}

void sigmoid(float* pd, int c, const float alpha, const float beta) {
  ORT_UNUSED_PARAMETER(alpha);
  ORT_UNUSED_PARAMETER(beta);

  clip_for_sigmoid_in_place(pd, c);

  for (int i = 0; i < c; i++) {
    float x = 0.5f * pd[i];
    float x2 = x * x;
    float p = x2 * alpha_13 + alpha_11;
    p = x2 * p + alpha_9;
    p = x2 * p + alpha_7;
    p = x2 * p + alpha_5;
    p = x2 * p + alpha_3;
    p = x2 * p + alpha_1;
    p = x * p;
    float q = x2 * beta_6 + beta_4;
    q = x2 * q + beta_2;
    q = x2 * q + beta_0;
    pd[i] = 0.5f * (1 + (p / q));
  }
}

void tanh(float* pd, int c, const float alpha, const float beta) {
  ORT_UNUSED_PARAMETER(alpha);
  ORT_UNUSED_PARAMETER(beta);

  clip_for_tanh_in_place(pd, c);

  for (int i = 0; i < c; i++) {
    float x = pd[i];
    float x2 = x * x;
    float p = x2 * alpha_13 + alpha_11;
    p = x2 * p + alpha_9;
    p = x2 * p + alpha_7;
    p = x2 * p + alpha_5;
    p = x2 * p + alpha_3;
    p = x2 * p + alpha_1;
    p = x * p;
    float q = x2 * beta_6 + beta_4;
    q = x2 * q + beta_2;
    q = x2 * q + beta_0;
    pd[i] = p / q;
  }
}

void relu(float* pd, int c, const float alpha, const float beta) {
  ORT_UNUSED_PARAMETER(alpha);
  ORT_UNUSED_PARAMETER(beta);

  for (int i = 0; i < c; i++) {
    if (pd[i] < 0)
      pd[i] = 0.0f;
  }
}

void sigmoid_exact(float* pd, int c, const float alpha, const float beta) {
  ORT_UNUSED_PARAMETER(alpha);
  ORT_UNUSED_PARAMETER(beta);

  for (int i = 0; i < c; i++) {
    float x = pd[i];
    pd[i] = 1.0f / (1 + ::std::exp(-x));
  }
}

void tanh_exact(float* pd, int c, const float alpha, const float beta) {
  ORT_UNUSED_PARAMETER(alpha);
  ORT_UNUSED_PARAMETER(beta);

  for (int i = 0; i < c; i++) {
    float x = pd[i];
    pd[i] = ::std::tanh(x);
  }
}

void merge_lstm_gates_to_memory(const float* pprev, const float* pi, const float* pf, const float* pg,
                                float* pcurr, const int c) {
  for (int i = 0; i < c; i++) {
    pcurr[i] = pprev[i] * pf[i] + pi[i] * pg[i];
  }
}

void gru_reset_gate_tanh(const float* ps1, float* ps2, float* pd, const int c,
                         const float alpha, const float beta) {
  ORT_UNUSED_PARAMETER(alpha);
  ORT_UNUSED_PARAMETER(beta);

  clip_for_tanh_in_place(ps2, c);

  for (int i = 0; i < c; i++) {
    float x = ps2[i];
    float x2 = x * x;
    float p = x2 * alpha_13 + alpha_11;
    p = x2 * p + alpha_9;
    p = x2 * p + alpha_7;
    p = x2 * p + alpha_5;
    p = x2 * p + alpha_3;
    p = x2 * p + alpha_1;
    p = x * p;
    float q = x2 * beta_6 + beta_4;
    q = x2 * q + beta_2;
    q = x2 * q + beta_0;
    pd[i] = ps1[i] * p / q;
  }
}

void gru_reset_gate_sigmoid(const float* ps1, float* ps2, float* pd, const int c,
                            const float alpha, const float beta) {
  ORT_UNUSED_PARAMETER(alpha);
  ORT_UNUSED_PARAMETER(beta);

  clip_for_sigmoid_in_place(ps2, c);

  for (int i = 0; i < c; i++) {
    float x = 0.5f * ps2[i];
    float x2 = x * x;
    float p = x2 * alpha_13 + alpha_11;
    p = x2 * p + alpha_9;
    p = x2 * p + alpha_7;
    p = x2 * p + alpha_5;
    p = x2 * p + alpha_3;
    p = x2 * p + alpha_1;
    p = x * p;
    float q = x2 * beta_6 + beta_4;
    q = x2 * q + beta_2;
    q = x2 * q + beta_0;
    pd[i] = ps1[i] * 0.5f * (1 + p / q);
  }
}

void gru_reset_gate_relu(const float* ps1, float* ps2, float* pd, const int c,
                         const float alpha, const float beta) {
  ORT_UNUSED_PARAMETER(alpha);
  ORT_UNUSED_PARAMETER(beta);

  for (int i = 0; i < c; i++) {
    const auto max = ps2[i] > 0 ? ps2[i] : 0.0f;
    pd[i] = ps1[i] * max;
  }
}

void gru_reset_gate_composed(const float* ps1, float* ps2, float* pd, const int c,
                             std::function<float(float, float, float)> func,
                             const float alpha, const float beta) {
  for (int i = 0; i < c; i++) {
    pd[i] = ps1[i] * func(ps2[i], alpha, beta);
  }
}

void gru_output_gate_tanh(float* ph, const float* pz, const float* ps, float* po, const int c,
                          const float alpha, const float beta) {
  ORT_UNUSED_PARAMETER(alpha);
  ORT_UNUSED_PARAMETER(beta);

  clip_for_tanh_in_place(ph, c);

  for (int i = 0; i < c; i++) {
    float x = ph[i];
    float x2 = x * x;
    float p = x2 * alpha_13 + alpha_11;
    p = x2 * p + alpha_9;
    p = x2 * p + alpha_7;
    p = x2 * p + alpha_5;
    p = x2 * p + alpha_3;
    p = x2 * p + alpha_1;
    p = x * p;
    float q = x2 * beta_6 + beta_4;
    q = x2 * q + beta_2;
    q = x2 * q + beta_0;
    po[i] = (1 - pz[i]) * (p / q) + pz[i] * ps[i];
  }
}

void gru_output_gate_relu(float* ph, const float* pz, const float* ps, float* po, const int c,
                          const float alpha, const float beta) {
  ORT_UNUSED_PARAMETER(alpha);
  ORT_UNUSED_PARAMETER(beta);

  for (int i = 0; i < c; i++) {
    float max = ph[i] > 0 ? ph[i] : 0.0f;
    po[i] = (1 - pz[i]) * max + pz[i] * ps[i];
  }
}

void gru_output_gate_composed(float* ph, const float* pz, const float* ps, float* po, const int c,
                              std::function<float(float, float, float)> func,
                              const float alpha, const float beta) {
  for (int i = 0; i < c; i++) {
    po[i] = (1 - pz[i]) * func(ph[i], alpha, beta) + pz[i] * ps[i];
  }
}

void gru_output_gate_sigmoid(float* ph, const float* pz, const float* ps, float* po, const int c,
                             const float alpha, const float beta) {
  ORT_UNUSED_PARAMETER(alpha);
  ORT_UNUSED_PARAMETER(beta);

  clip_for_sigmoid_in_place(ph, c);

  for (int i = 0; i < c; i++) {
    float x = 0.5f * (ph[i]);
    float x2 = x * x;
    float p = x2 * alpha_13 + alpha_11;
    p = x2 * p + alpha_9;
    p = x2 * p + alpha_7;
    p = x2 * p + alpha_5;
    p = x2 * p + alpha_3;
    p = x2 * p + alpha_1;
    p = x * p;
    float q = x2 * beta_6 + beta_4;
    q = x2 * q + beta_2;
    q = x2 * q + beta_0;
    po[i] = (1 - pz[i]) * 0.5f * (1 + (p / q)) + pz[i] * ps[i];
  }
}

void composed_activation_func(float* ps, const int c,
                              std::function<float(float, float, float)> func,
                              const float alpha, const float beta) {
  for (int i = 0; i < c; i++) {
    ps[i] = func(ps[i], alpha, beta);
  }
}

void composed_lstm_merge_gates_func(float* ps, const int c,
                                    std::function<float(float, float, float)> func,
                                    const float alpha, const float beta) {
  for (int i = 0; i < c; i++) {
    ps[i] = func(ps[i], alpha, beta);
  }
}

void composed_gru_reset_gate_func(float* ps, const int c,
                                  std::function<float(float, float, float)> func,
                                  const float alpha, const float beta) {
  for (int i = 0; i < c; i++) {
    ps[i] = func(ps[i], alpha, beta);
  }
}

void composed_gru_output_gate_func(float* ps, const int c,
                                   std::function<float(float, float, float)> func,
                                   const float alpha, const float beta) {
  for (int i = 0; i < c; i++) {
    ps[i] = func(ps[i], alpha, beta);
  }
}

ActivationFuncPtr ActivationFuncByName(const std::string& func) {
  if (func == "sigmoid")
    return sigmoid;

  if (func == "tanh")
    return tanh;

  if (func == "relu")
    return relu;

  if (func == "affine")
    return [](float* ps, const int c, const float alpha, const float beta) {
      composed_activation_func(ps, c, Affine<float>, alpha, beta);
    };

  if (func == "leakyrelu")
    return [](float* ps, const int c, const float alpha, const float beta) {
      composed_activation_func(ps, c, LeakyRelu<float>, alpha, beta);
    };

  if (func == "thresholdedrelu")
    return [](float* ps, const int c, const float alpha, const float beta) {
      composed_activation_func(ps, c, ThresholdedRelu<float>, alpha, beta);
    };

  if (func == "scaledtanh")
    return [](float* ps, const int c, const float alpha, const float beta) {
      composed_activation_func(ps, c, ScaledTanh<float>, alpha, beta);
    };

  if (func == "hardsigmoid")
    return [](float* ps, const int c, const float alpha, const float beta) {
      composed_activation_func(ps, c, HardSigmoid<float>, alpha, beta);
    };

  if (func == "elu")
    return [](float* ps, const int c, const float alpha, const float beta) {
      composed_activation_func(ps, c, Elu<float>, alpha, beta);
    };

  if (func == "softsign")
    return [](float* ps, const int c, const float alpha, const float beta) {
      composed_activation_func(ps, c, Softsign<float>, alpha, beta);
    };

  if (func == "softplus")
    return [](float* ps, const int c, const float alpha, const float beta) {
      composed_activation_func(ps, c, Softplus<float>, alpha, beta);
    };

  ORT_THROW("Invalid activation function of ", func);
}

LstmMergeGatesFuncPtr LstmMergeGatesFuncByName(const std::string& func) {
  if (func == "sigmoid")
    return sigmoid_m;

  if (func == "tanh")
    return tanh_m;

  if (func == "relu")
    return relu_m;

  if (func == "affine")
    return [](const float* ps1, float* ps1_c, const float* ps2, float* ps3, const int c,
              const float alpha, const float beta) {
      composed_m(ps1, ps1_c, ps2, ps3, c, Affine<float>, alpha, beta);
    };

  if (func == "leakyrelu")
    return [](const float* ps1, float* ps1_c, const float* ps2, float* ps3, const int c,
              const float alpha, const float beta) {
      composed_m(ps1, ps1_c, ps2, ps3, c, LeakyRelu<float>, alpha, beta);
    };

  if (func == "thresholdedrelu")
    return [](const float* ps1, float* ps1_c, const float* ps2, float* ps3, const int c,
              const float alpha, const float beta) {
      composed_m(ps1, ps1_c, ps2, ps3, c, ThresholdedRelu<float>, alpha, beta);
    };

  if (func == "scaledtanh")
    return [](const float* ps1, float* ps1_c, const float* ps2, float* ps3, const int c,
              const float alpha, const float beta) {
      composed_m(ps1, ps1_c, ps2, ps3, c, ScaledTanh<float>, alpha, beta);
    };

  if (func == "hardsigmoid")
    return [](const float* ps1, float* ps1_c, const float* ps2, float* ps3, const int c,
              const float alpha, const float beta) {
      composed_m(ps1, ps1_c, ps2, ps3, c, HardSigmoid<float>, alpha, beta);
    };

  if (func == "elu")
    return [](const float* ps1, float* ps1_c, const float* ps2, float* ps3, const int c,
              const float alpha, const float beta) {
      composed_m(ps1, ps1_c, ps2, ps3, c, Elu<float>, alpha, beta);
    };

  if (func == "softsign")
    return [](const float* ps1, float* ps1_c, const float* ps2, float* ps3, const int c,
              const float alpha, const float beta) {
      composed_m(ps1, ps1_c, ps2, ps3, c, Softsign<float>, alpha, beta);
    };

  if (func == "softplus")
    return [](const float* ps1, float* ps1_c, const float* ps2, float* ps3, const int c,
              const float alpha, const float beta) {
      composed_m(ps1, ps1_c, ps2, ps3, c, Softplus<float>, alpha, beta);
    };

  ORT_THROW("Invalid LSTM merge activation function of ", func);
}

GruResetGateFuncPtr GruResetGateFuncByName(const std::string& func) {
  if (func == "sigmoid")
    return gru_reset_gate_sigmoid;

  if (func == "tanh")
    return gru_reset_gate_tanh;

  if (func == "relu")
    return gru_reset_gate_relu;

  if (func == "affine")
    return [](const float* ps1, float* ps2, float* ps3, const int c, const float alpha, const float beta) {
      gru_reset_gate_composed(ps1, ps2, ps3, c, Affine<float>, alpha, beta);
    };

  if (func == "leakyrelu")
    return [](const float* ps1, float* ps2, float* ps3, const int c, const float alpha, const float beta) {
      gru_reset_gate_composed(ps1, ps2, ps3, c, LeakyRelu<float>, alpha, beta);
    };

  if (func == "thresholdedrelu")
    return [](const float* ps1, float* ps2, float* ps3, const int c, const float alpha, const float beta) {
      gru_reset_gate_composed(ps1, ps2, ps3, c, ThresholdedRelu<float>, alpha, beta);
    };

  if (func == "scaledtanh")
    return [](const float* ps1, float* ps2, float* ps3, const int c, const float alpha, const float beta) {
      gru_reset_gate_composed(ps1, ps2, ps3, c, ScaledTanh<float>, alpha, beta);
    };

  if (func == "hardsigmoid")
    return [](const float* ps1, float* ps2, float* ps3, const int c, const float alpha, const float beta) {
      gru_reset_gate_composed(ps1, ps2, ps3, c, HardSigmoid<float>, alpha, beta);
    };

  if (func == "elu")
    return [](const float* ps1, float* ps2, float* ps3, const int c, const float alpha, const float beta) {
      gru_reset_gate_composed(ps1, ps2, ps3, c, Elu<float>, alpha, beta);
    };

  if (func == "softsign")
    return [](const float* ps1, float* ps2, float* ps3, const int c, const float alpha, const float beta) {
      gru_reset_gate_composed(ps1, ps2, ps3, c, Softsign<float>, alpha, beta);
    };

  if (func == "softplus")
    return [](const float* ps1, float* ps2, float* ps3, const int c, const float alpha, const float beta) {
      gru_reset_gate_composed(ps1, ps2, ps3, c, Softplus<float>, alpha, beta);
    };

  ORT_THROW("Invalid GRU reset gate activation function: ", func);
}

GruOutputGateFuncPtr GruOutputGateFuncByName(const std::string& func) {
  if (func == "sigmoid")
    return gru_output_gate_sigmoid;

  if (func == "tanh")
    return gru_output_gate_tanh;

  if (func == "relu")
    return gru_output_gate_relu;

  if (func == "affine")
    return [](float* ps1, const float* ps2, const float* ph, float* ps3, const int c,
              const float alpha, const float beta) {
      gru_output_gate_composed(ps1, ps2, ph, ps3, c, Affine<float>, alpha, beta);
    };

  if (func == "leakyrelu")
    return [](float* ps1, const float* ps2, const float* ph, float* ps3, const int c,
              const float alpha, const float beta) {
      gru_output_gate_composed(ps1, ps2, ph, ps3, c, LeakyRelu<float>, alpha, beta);
    };

  if (func == "thresholdedrelu")
    return [](float* ps1, const float* ps2, const float* ph, float* ps3, const int c,
              const float alpha, const float beta) {
      gru_output_gate_composed(ps1, ps2, ph, ps3, c, ThresholdedRelu<float>, alpha, beta);
    };

  if (func == "scaledtanh")
    return [](float* ps1, const float* ps2, const float* ph, float* ps3, const int c,
              const float alpha, const float beta) {
      gru_output_gate_composed(ps1, ps2, ph, ps3, c, ScaledTanh<float>, alpha, beta);
    };

  if (func == "hardsigmoid")
    return [](float* ps1, const float* ps2, const float* ph, float* ps3, const int c,
              const float alpha, const float beta) {
      gru_output_gate_composed(ps1, ps2, ph, ps3, c, HardSigmoid<float>, alpha, beta);
    };

  if (func == "elu")
    return [](float* ps1, const float* ps2, const float* ph, float* ps3, const int c,
              const float alpha, const float beta) {
      gru_output_gate_composed(ps1, ps2, ph, ps3, c, Elu<float>, alpha, beta);
    };

  if (func == "softsign")
    return [](float* ps1, const float* ps2, const float* ph, float* ps3, const int c,
              const float alpha, const float beta) {
      gru_output_gate_composed(ps1, ps2, ph, ps3, c, Softsign<float>, alpha, beta);
    };

  if (func == "softplus")
    return [](float* ps1, const float* ps2, const float* ph, float* ps3, const int c,
              const float alpha, const float beta) {
      gru_output_gate_composed(ps1, ps2, ph, ps3, c, Softplus<float>, alpha, beta);
    };

  ORT_THROW("Invalid GRU hidden gate activation function: ", func);
}

}  // namespace deepcpu
}  // namespace detail
}  // namespace rnn
}  // namespace onnxruntime
