Change order.

Author: trivialfis

src/common/stats.cc

@@ -45,15 +45,14 @@ void Median(Context const* ctx, linalg::Matrix<float> const& t,
   }
 }
 
-void Mean(Context const* ctx, linalg::Vector<float> const& v, linalg::Vector<float>* out) {
-  v.SetDevice(ctx->Device());
+void Mean(Context const* ctx, linalg::VectorView<float const> v, linalg::Vector<float>* out) {
   out->SetDevice(ctx->Device());
   out->Reshape(1);
 
   if (ctx->IsCUDA()) {
-    cuda_impl::Mean(ctx, v.View(ctx->Device()), out->View(ctx->Device()));
+    cuda_impl::Mean(ctx, v, out->View(ctx->Device()));
   } else {
-    auto h_v = v.HostView();
+    auto h_v = v;
     float n = v.Size();
     MemStackAllocator<float, DefaultMaxThreads()> tloc(ctx->Threads(), 0.0f);
     ParallelFor(v.Size(), ctx->Threads(),

src/common/stats.h

@@ -149,7 +149,7 @@ void Median(Context const* ctx, linalg::Matrix<float> const& t,
 /**
  * @brief Calculate the mean value of a vector.
  */
-void Mean(Context const* ctx, linalg::Vector<float> const& v, linalg::Vector<float>* out);
+void Mean(Context const* ctx, linalg::VectorView<float const> v, linalg::Vector<float>* out);
 
 /**
  * @brief Calculate the mean value for the first axis.

src/objective/multiclass_obj.cu

@@ -210,14 +210,15 @@ class SoftmaxMultiClassObj : public ObjFunction {
     collective::SafeColl(status);
     CHECK_GE(sum_weight, kRtEps);
     linalg::VecScaDiv(this->ctx_, intercept, sum_weight);
+    CHECK_EQ(base_score->Size(), n_classes);
 
     // Transform it back to margin
-    linalg::Vector<float> mean_intercepts;
-    CHECK_EQ(base_score->Size(), n_classes);
-    common::Mean(this->ctx_, *base_score, &mean_intercepts);
-    auto d_mean = mean_intercepts.View(this->ctx_->Device());
-    TransformKernel(this->ctx_, intercept,
-                    [=] XGBOOST_DEVICE(float v) { return log(v) - d_mean(0); });
+    // ln(v) - E[ln(v)]
+    linalg::Vector<float> mean;
+    linalg::LogE(this->ctx_, intercept);
+    common::Mean(this->ctx_, intercept, &mean);
+    auto d_mean = mean.View(this->ctx_->Device());
+    TransformKernel(this->ctx_, intercept, [=] XGBOOST_DEVICE(float v) { return v - d_mean(0); });
   }
 
  private:

tests/cpp/common/test_stats.cc

@@ -114,7 +114,7 @@ void TestMean(Context const* ctx) {
   float mean = nf * (nf - 1) / 2 / n;
 
   linalg::Vector<float> res{{1}, ctx->Device()};
-  Mean(ctx, data, &res);
+  Mean(ctx, data.View(ctx->Device()), &res);
   auto h_res = res.HostView();
   ASSERT_EQ(h_res.Size(), 1);
   ASSERT_EQ(mean, h_res(0));