inline void Softmax(const SoftmaxParams& params,
                    const RuntimeShape& input_shape, const In* input_data,
                    const RuntimeShape& output_shape, Out* output_data) {
  const int trailing_dim = input_shape.DimensionsCount() - 1;
  const int excluding_last_dim =
      MatchingFlatSizeSkipDim(input_shape, trailing_dim, output_shape);
  const int last_dim =
      MatchingDim(input_shape, trailing_dim, output_shape, trailing_dim);

  const int32_t clamp_max = std::numeric_limits<Out>::max();
  const int32_t clamp_min = std::numeric_limits<Out>::min();
  for (int i = 0; i < excluding_last_dim; ++i) {
    int32_t max_val = std::numeric_limits<In>::min();
    // Find max quantized value.
    for (int j = 0; j < last_dim; ++j) {
      max_val = std::max(max_val, static_cast<int32_t>(input_data[j]));
    }

    float sum_exp = 0.0f;
    const int32_t max_uint8 = std::numeric_limits<uint8_t>::max();
    const float* table_offset = &params.table[max_uint8 - max_val];
    // Calculate normalizer sum(exp(x)).
    for (int j = 0; j < last_dim; ++j) {
      sum_exp += table_offset[input_data[j]];
    }

    const float inv_sum_exp = 1.0f / (sum_exp * params.scale);
    // Normalize and quantize probabilities.
    for (int j = 0; j < last_dim; ++j) {
      const float prob_rescaled = table_offset[input_data[j]] * inv_sum_exp;
      const int32_t prob_quantized =
          QuantizeSoftmaxOutput<Out>(prob_rescaled, params.zero_point);
      output_data[j] = static_cast<Out>(
          std::max(std::min(clamp_max, prob_quantized), clamp_min));
    }
    input_data += last_dim;
    output_data += last_dim;
  }
}
inline void Softmax(const SoftmaxParams& params,
                    const RuntimeShape& input_shape, const In* input_data,
                    const RuntimeShape& output_shape, Out* output_data) {
  const int trailing_dim = input_shape.DimensionsCount() - 1;
  const int excluding_last_dim =
      MatchingFlatSizeSkipDim(input_shape, trailing_dim, output_shape);
  const int last_dim =
      MatchingDim(input_shape, trailing_dim, output_shape, trailing_dim);

  const int32_t clamp_max = std::numeric_limits<Out>::max();
  const int32_t clamp_min = std::numeric_limits<Out>::min();
  for (int i = 0; i < excluding_last_dim; ++i) {
    int32_t max_val = std::numeric_limits<In>::min();
     
    for (int j = 0; j < last_dim; ++j) {
      max_val = std::max(max_val, static_cast<int32_t>(input_data[j]));
    }

    float sum_exp = 0.0f;
    const int32_t max_uint8 = std::numeric_limits<uint8_t>::max();
    const float* table_offset = &params.table[max_uint8 - max_val];
     
    for (int j = 0; j < last_dim; ++j) {
      sum_exp += table_offset[input_data[j]];
    }

    const float inv_sum_exp = 1.0f / (sum_exp * params.scale);
     
    for (int j = 0; j < last_dim; ++j) {
      const float prob_rescaled = table_offset[input_data[j]] * inv_sum_exp;
      const int32_t prob_quantized =
          QuantizeSoftmaxOutput<Out>(prob_rescaled, params.zero_point);
      output_data[j] = static_cast<Out>(
          std::max(std::min(clamp_max, prob_quantized), clamp_min));
    }
    input_data += last_dim;
    output_data += last_dim;
  }
}