chore: import upstream snapshot with attribution

paddle/phi/kernels/cpu/nce_kernel.cc

@@ -0,0 +1,246 @@
+// Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include <math.h>
+
+#include <iterator>
+#include <random>
+
+#include "paddle/phi/core/kernel_registry.h"
+#include "paddle/phi/kernels/funcs/eigen/common.h"
+#include "paddle/phi/kernels/funcs/math/sampler.h"
+#include "paddle/utils/optional.h"
+
+namespace phi {
+
+using Sampler = math::Sampler;
+
+template <typename Context, typename T>
+static void inline PrepareSamples(const Context &dev_ctx,
+                                  Sampler *sampler,
+                                  DenseTensor *sample_labels,
+                                  const DenseTensor &label_in,
+                                  const std::vector<int> &custom_neg_classes) {
+  auto label = &label_in;
+  const int64_t *label_data = label->data<int64_t>();
+  auto label_dims = label->dims();
+
+  auto sample_labels_dims = sample_labels->dims();
+  int64_t *sample_labels_data = dev_ctx.template Alloc<int64_t>(sample_labels);
+
+  int64_t num_label = label_dims.size() == 2 ? label_dims[1] : 1;
+  int64_t index = 0;
+  for (int64_t i = 0; i < label_dims[0]; ++i) {
+    int64_t j = 0;
+    for (; j < num_label; ++j) {
+      sample_labels_data[index++] = label_data[i * num_label + j];
+    }
+    // for unittest
+    if (custom_neg_classes.size() > 0) {
+      for (auto label : custom_neg_classes) {
+        sample_labels_data[index++] = label;
+      }
+    } else {
+      for (; j < sample_labels_dims[1]; ++j) {
+        // TODO(wanghaoshuang): support more distribution sampling
+        sample_labels_data[index++] = sampler->Sample();
+      }
+    }
+  }
+}
+
+template <typename T, typename Context>
+void NCEKernel(const Context &dev_ctx,
+               const DenseTensor &input_in,
+               const DenseTensor &label_in,
+               const DenseTensor &weight_in,
+               const optional<DenseTensor> &bias_in,
+               const optional<DenseTensor> &sample_weight_in,
+               const optional<DenseTensor> &custom_dist_probs,
+               const optional<DenseTensor> &custom_dist_alias,
+               const optional<DenseTensor> &custom_dist_alias_probs,
+               int num_total_classes,
+               const std::vector<int> &custom_neg_classes,
+               int num_neg_samples,
+               int sampler_in,
+               int seed,
+               bool is_sparse,
+               bool remote_prefetch,
+               bool is_test,
+               DenseTensor *cost_out,
+               DenseTensor *sample_logits_out,
+               DenseTensor *sample_labels_out) {
+  int sampler_type = sampler_in;
+
+  Sampler *sampler;
+  switch (sampler_type) {
+    case 0: {
+      sampler = new math::UniformSampler(num_total_classes - 1, seed);
+      break;
+    }
+    case 1: {
+      sampler = new math::LogUniformSampler(num_total_classes - 1, seed);
+      break;
+    }
+    case 2: {
+      auto dist_probs = custom_dist_probs.get_ptr();
+      auto dist_alias = custom_dist_alias.get_ptr();
+      auto dist_alias_probs = custom_dist_alias_probs.get_ptr();
+
+      PADDLE_ENFORCE_EQ(
+          dist_probs->numel(),
+          num_total_classes,
+          common::errors::InvalidArgument(
+              "ShapeError: The number of elements in Input(CustomDistProbs) "
+              "should be equal to the number of total classes. But Received: "
+              "Input(CustomDistProbs).numel() = %d, Attr(num_total_classes) "
+              "= %d.",
+              dist_probs->numel(),
+              num_total_classes));
+      PADDLE_ENFORCE_EQ(
+          dist_alias->numel(),
+          num_total_classes,
+          common::errors::InvalidArgument(
+              "ShapeError: The number of elements in Input(CustomDistAlias) "
+              "should be equal to the number of total classes. But Received: "
+              "Input(CustomDistAlias).numel() = %d, Attr(num_total_classes) "
+              "= %d.",
+              dist_alias->numel(),
+              num_total_classes));
+      PADDLE_ENFORCE_EQ(
+          dist_alias_probs->numel(),
+          num_total_classes,
+          common::errors::InvalidArgument(
+              "ShapeError: The number of elements in "
+              "Input(CustomDistAliasProbs) "
+              "should be equal to the number of total classes. But Received: "
+              "Input(CustomDistAliasProbs).numel() = %d, "
+              "Attr(num_total_classes) = %d.",
+              dist_alias_probs->numel(),
+              num_total_classes));
+
+      const float *probs_data = dist_probs->data<float>();
+      const int *alias_data = dist_alias->data<int>();
+      const float *alias_probs_data = dist_alias_probs->data<float>();
+      sampler = new math::CustomSampler(num_total_classes - 1,
+                                        probs_data,
+                                        alias_data,
+                                        alias_probs_data,
+                                        seed);
+      break;
+    }
+    default: {
+      PADDLE_THROW(common::errors::InvalidArgument(
+          "Unsupported SamplerType. SamplerType should be 0: Uniform, "
+          "1: LogUniform or 2: CustomDist. Received SamplerType: %d",
+          sampler_type));
+    }
+  }
+
+  std::vector<int64_t> sample_out_dims;
+  auto label = &label_in;
+  DenseTensor *sample_labels;
+  DenseTensor *sample_out;
+  DenseTensor sample_labels_tmp, sample_out_tmp;
+  if (is_test) {
+    // set dims of output(SampleOut)
+    int64_t num_true_classes = label->dims().size() == 2 ? label->dims()[1] : 1;
+    sample_out_dims.push_back(input_in.dims()[0]);
+    sample_out_dims.push_back(
+        (num_true_classes == -1) ? -1 : (num_neg_samples + num_true_classes));
+
+    sample_labels = &sample_labels_tmp;
+    sample_labels->Resize(sample_out_dims);
+
+    sample_out = &sample_out_tmp;
+    sample_out->Resize(sample_out_dims);
+  } else {
+    sample_labels = sample_labels_out;
+    sample_out = sample_logits_out;
+  }
+
+  PrepareSamples<Context, T>(
+      dev_ctx, sampler, sample_labels, label_in, custom_neg_classes);
+  const int64_t *sample_labels_data = sample_labels->data<int64_t>();
+
+  for (int64_t x = 0; x < sample_labels->numel(); x++) {
+    PADDLE_ENFORCE_GE(sample_labels_data[x],
+                      0,
+                      common::errors::InvalidArgument(
+                          "ValueError: Every sample label should be "
+                          "non-negative. But received: "
+                          "Input(SampleLabels)[%d] = %d",
+                          x,
+                          sample_labels_data[x]));
+  }
+
+  T *sample_out_data = dev_ctx.template Alloc<T>(sample_out);
+  auto sample_weight = sample_weight_in.get_ptr();
+  const T *sample_weight_data = nullptr;
+  if (sample_weight != nullptr) {
+    sample_weight_data = sample_weight->data<T>();
+  }
+  auto out = cost_out;
+  T *out_data = dev_ctx.template Alloc<T>(out);
+  int64_t num_true_class = 1;
+  if (label != nullptr) {
+    num_true_class = label->dims()[1];
+  }
+  int64_t sampled_labels_num = sample_labels->dims()[1];
+  //    T b = 1. / num_total_classes * num_neg_samples;
+  // forward bias
+  auto bias = bias_in.get_ptr();
+  if (bias != nullptr) {
+    const T *bias_data = bias->data<T>();
+    for (int64_t i = 0; i < sample_labels->numel(); ++i) {
+      sample_out_data[i] = bias_data[sample_labels_data[i]];
+    }
+  } else {
+    for (int64_t i = 0; i < sample_labels->numel(); ++i) {
+      sample_out_data[i] = 0;
+    }
+  }
+  // forward mul
+  auto input_mat = EigenMatrix<T>::From(input_in);
+
+  auto weight_mat = EigenMatrix<T>::From(weight_in);
+  for (int64_t i = 0; i < sample_labels->numel(); ++i) {
+    Eigen::Tensor<T, 0, Eigen::RowMajor, int64_t> result =
+        (input_mat.chip(static_cast<int>(i / sample_labels->dims()[1]), 0) *
+         weight_mat.chip(sample_labels_data[i], 0))
+            .sum();
+    sample_out_data[i] += result(0);
+    sample_out_data[i] = (1. / (1. + exp(-sample_out_data[i])));
+  }
+
+  // forward cost
+  for (int64_t i = 0; i < sample_labels->dims()[0]; ++i) {
+    out_data[i] = 0;
+    T w = sample_weight == nullptr ? 1. : sample_weight_data[i];
+    for (int64_t j = 0; j < sampled_labels_num; ++j) {
+      int64_t target = sample_labels_data[i * sampled_labels_num + j];
+      T o = sample_out_data[i * sampled_labels_num + j];
+      float b = sampler->Probability(target) * num_neg_samples;
+      T cost = (j < num_true_class) ? -log(o / (o + b)) : -log(b / (o + b));
+      out_data[i] += w * cost;
+    }
+  }
+  delete sampler;
+}
+
+}  // namespace phi
+
+PD_REGISTER_KERNEL(nce, CPU, ALL_LAYOUT, phi::NCEKernel, float, double) {
+  kernel->OutputAt(2).SetDataType(phi::DataType::INT64);
+}