Align diopiStd with camb, torch_cuda, ascend for >=torch2.0

diopi_test/python/configs/diopi_configs.py

@@ -3116,7 +3116,7 @@
         interface=['torch'],
         para=dict(
             dim=[-1, 0, 1, [0, 1], 2, [-1, 0, -3], [0, 2, 3, -1], -1, [-1, -2], 2, None, None],
-            unbiased=[True, True, False, True, False, True, False, False, True, False, False, True],
+            correction=[1, 1, 0, 1, 0, 1, 0, 0, 1, 0, 0, 1],
         ),
         atol=1e-4,
         rtol=1e-5,

diopi_test/python/conformance/diopi_functions.py

@@ -808,15 +808,16 @@ def mean(input, dim=None, keepdim=False, dtype=None) -> Tensor:
     return out
 
 
-def std(input, unbiased=True, dim=None, keepdim=False) -> Tensor:
+def std(input, correction=1, dim=None, keepdim=False) -> Tensor:
     assert (
         isinstance(dim, (int, list)) or dim is None
     ), "dim should be int or list or None"
 
     dim, out = reduce_op_process(input, dim, keepdim)
     dim1 = Sizes(list(dim))
     func = check_function("diopiStd")
-    ret = func(input.context(), out, input, dim1, unbiased)
+    correction = Scalar(correction)
+    ret = func(input.context(), out, input, dim1, correction)
     check_returncode(ret)
     return out
 

diopi_test/python/docs/CN_doc/__init__.py

@@ -534,11 +534,10 @@ def add_docstr(attr, docstr):
 
 add_docstr("std", r"""
     释义
-        如果 *unbiased* 为 ``True``, 则将使用 *Bessel* 校正。 否则, 将直接计算样本偏差, 而不进行任何校正。
         如果 *dim* 等于 ``None``, 将对所有元素计算标准差。
     参数
         - *input* ( **Tensor** ) : 输入张量
-        - *unbiased* ( **bool** ) : 是否使用Bessel校正
+        - *correction* ( **int** ) : Bessel校正的修正值
         - *dim* ( **int** 或者 **list(int)** ) : 进行归约的维度
         - *keepdim* ( **bool** ) : 结果是否保留原有维度
     C API

diopi_test/python/docs/EN_doc/__init__.py

@@ -1150,21 +1150,19 @@ def add_docstr(attr, docstr):
 add_docstr("std", r"""
 Calculates the standard deviation of each row of attr:`input` tensor in the given
 dimension :attr:`dim`.
-If :attr:`unbiased` is ``True``, Bessel's correction will be used.
-Otherwise, the sample deviation is calculated, without any correction.
 
 Args:
     {input}
     {dim}
 
 Keyword args:
-    unbiased (bool): whether to use Bessel's correction (:math:`\delta N = 1`).
+    correction (int): difference between the sample size and sample degrees of freedom. Defaults to Bessel’s correction. (:math:`\delta N = 1`).
     {keepdim}
 
 Example::
 
     >>> a = tensor([[-0.8166, -1.3802, -0.3560]])
-    >>> std(a, unbiased=False)
+    >>> std(a, correction=0)
     tensor(0.4188)
 """.format(**multi_dim_common))
 

impl/ascend/functions/reduce.cpp

@@ -54,7 +54,7 @@ diopiError_t diopiMean(diopiContextHandle_t ctx, diopiTensorHandle_t out, diopiC
     return diopiSuccess;
 }
 
-diopiError_t diopiStd(diopiContextHandle_t ctx, diopiTensorHandle_t out, diopiConstTensorHandle_t input, diopiSize_t dim, bool unbiased) {
+diopiError_t diopiStd(diopiContextHandle_t ctx, diopiTensorHandle_t out, diopiConstTensorHandle_t input, diopiSize_t dim, const diopiScalar_t* correction) {
     AscendTensor inputAt(input);
     AscendTensor outAt(out);
 
@@ -63,18 +63,18 @@ diopiError_t diopiStd(diopiContextHandle_t ctx, diopiTensorHandle_t out, diopiCo
         keepdim = true;
     }
 
-    int64_t correction = 0;
-    if (unbiased) {
-        correction = 1;
+    int64_t correctionInt = 1;
+    if (correction != nullptr) {
+        correctionInt = correction->ival;
     }
 
     if (dim.data == nullptr || dim.len == 0) {
         std::vector<int64_t> allDim(inputAt.dim());
         std::iota(allDim.begin(), allDim.end(), 0);
         diopiSize_t rDim = vectorToDiopiSize(allDim);
-        DIOPI_ASCEND_CALL_ACLNN(aclnnStd, ctx, input, rDim, correction, keepdim, out);
+        DIOPI_ASCEND_CALL_ACLNN(aclnnStd, ctx, input, rDim, correctionInt, keepdim, out);
     } else {
-        DIOPI_ASCEND_CALL_ACLNN(aclnnStd, ctx, input, dim, correction, keepdim, out);
+        DIOPI_ASCEND_CALL_ACLNN(aclnnStd, ctx, input, dim, correctionInt, keepdim, out);
     }
     return diopiSuccess;
 }

impl/ascend_npu/diopi_impl/std.cpp

@@ -11,21 +11,23 @@
 
 namespace OP_IMPL_NS {
 
-diopiError_t diopiStd(diopiContextHandle_t ctx, diopiTensorHandle_t out, diopiConstTensorHandle_t input, diopiSize_t dim, bool unbiased) {
-    BEGIN_CALL_ACL_OP(out, input, dim);
+diopiError_t diopiStd(diopiContextHandle_t ctx, diopiTensorHandle_t out, diopiConstTensorHandle_t input, diopiSize_t dim, const diopiScalar_t* correction) {
+    BEGIN_CALL_ACL_OP(out, input, dim, correction);
+    if (correction == nullptr) {
+        correctionAt = 1;  // default correction value in torch_std is 1
+    }
     bool keepdim = false;
     if (inputAt.dim() == outAt.dim()) {
         keepdim = true;
     }
-    at::Scalar correction(static_cast<int64_t>(unbiased));
     if (0 == dim.len) {
         c10::DimVector adim(inputAt.dim());
         std::iota(adim.begin(), adim.end(), 0);
         at::IntArrayRef rdim(adim.data(), adim.size());
-        op_api::std_out(inputAt, rdim, correction, keepdim, outAt);
+        op_api::std_out(inputAt, rdim, correctionAt, keepdim, outAt);
     } else {
         at::IntArrayRef rdim(dim.data, dim.len);
-        op_api::std_out(inputAt, rdim, correction, keepdim, outAt);
+        op_api::std_out(inputAt, rdim, correctionAt, keepdim, outAt);
     }
     return diopiSuccess;
 }

impl/camb/functions/std.cpp

@@ -23,21 +23,27 @@ namespace camb {
  * @param unbiased whether to compute the unbiased standard deviation.
  * @param[out] out the output tensor depend on dim. type = [float32, float64, float16].
  */
-DIOPI_API diopiError_t diopiStd(diopiContextHandle_t ctx, diopiTensorHandle_t out, diopiConstTensorHandle_t input, diopiSize_t dim, bool unbiased) {
+DIOPI_API diopiError_t diopiStd(diopiContextHandle_t ctx, diopiTensorHandle_t out, diopiConstTensorHandle_t input, diopiSize_t dim,
+                                const diopiScalar_t* correction) {
     cnnlHandle_t handle = cnnlHandlePool.get(ctx);
 
     // shape of outTensor does not keep the dim.
     DiopiTensor outTensor(out);
     DiopiTensor inputTensor(input);
     auto outDtype = outTensor.dtype();
 
+    bool unbiased = true;
+    if (correction != nullptr) {
+        unbiased = correction->ival;
+    }
+
     bool keepDim = false;
     if (outTensor.dim() == inputTensor.dim()) {
         keepDim = true;
     }
 
     int axisNum = 0;
-    int *axis = nullptr;
+    int* axis = nullptr;
     if (0 == dim.len) {
         axisNum = inputTensor.dim();
         axis = new int[axisNum];
@@ -65,7 +71,7 @@ DIOPI_API diopiError_t diopiStd(diopiContextHandle_t ctx, diopiTensorHandle_t ou
     }
 
     // cast supported dtyeps for tensors.
-    std::vector<DiopiTensor *> tensorVecPtr{&outTensor, &inputTensor};
+    std::vector<DiopiTensor*> tensorVecPtr{&outTensor, &inputTensor};
     std::set<diopiDtype_t> supportedDtype{diopi_dtype_float16, diopi_dtype_float32, diopi_dtype_float64};
     DIOPI_CALL(autoCastTensorType(ctx, tensorVecPtr, supportedDtype));
     outTensor = *tensorVecPtr[0];
@@ -81,7 +87,7 @@ DIOPI_API diopiError_t diopiStd(diopiContextHandle_t ctx, diopiTensorHandle_t ou
     delete[] axis;
 
     size_t workspaceSize = 0;
-    void *workspace = nullptr;
+    void* workspace = nullptr;
     DIOPI_CALL_CNNL(cnnlGetStdVarMeanWorkspaceSize(handle, stdVarMeanDesc, inputDesc.get(), &workspaceSize));
     if (workspaceSize > 0) {
         workspace = requiresBuffer(ctx, workspaceSize).data();

impl/torch/functions/functions.cpp

@@ -389,17 +389,30 @@ diopiError_t diopiSum(diopiContextHandle_t ctx, diopiTensorHandle_t out, diopiCo
     return diopiSuccess;
 }
 
-diopiError_t diopiStd(diopiContextHandle_t ctx, diopiTensorHandle_t out, diopiConstTensorHandle_t input, diopiSize_t dim, bool unbiased) {
+diopiError_t diopiStd(diopiContextHandle_t ctx, diopiTensorHandle_t out, diopiConstTensorHandle_t input, diopiSize_t dim, const diopiScalar_t* correction) {
     impl::aten::setCurStream(ctx);
     auto atInput = impl::aten::buildATen(input);
     auto atOut = impl::aten::buildATen(out);
     auto atDim = impl::aten::buildAtIntArray(dim);
+
     bool keepdim = false;
     if (atInput.dim() == atOut.dim()) {
         keepdim = true;
     }
-    CALL_ATEN_CUDA_FUNC(std_out, atOut, atInput, atDim, unbiased, keepdim);
 
+#if TORCH_MM_VERSION >= 2010
+    c10::optional<at::Scalar> atCorrection = c10::optional<at::Scalar>();
+    if (correction != nullptr) {
+        atCorrection = impl::aten::buildAtScalar(correction);
+    }
+#else
+    c10::optional<int64_t> atCorrection = c10::optional<int64_t>();
+    if (correction != nullptr) {
+        atCorrection = impl::aten::buildAtScalar(correction).toInt();
+    }
+#endif
+
+    CALL_ATEN_CUDA_FUNC(std_out, atOut, atInput, atDim, atCorrection, keepdim);
     return diopiSuccess;
 }
 

proto/include/diopi/functions.h

@@ -1860,10 +1860,11 @@ DIOPI_API diopiError_t diopiSum(diopiContextHandle_t ctx, diopiTensorHandle_t ou
  * @param[in] ctx Context environment.
  * @param[in] input the input tensor, type = [float32, float64, float16].
  * @param[in] dim an array, dimension for reduction. type = [int32, int64].
- * @param[in] unbiased whether to compute the unbiased standard deviation.
+ * @param[in] correction difference between the sample size and sample degrees of freedom. Defaults to Bessel’s correction.
  * @param[out] out the output tensor depend on dim. type = [float32, float64, float16].
  */
-DIOPI_API diopiError_t diopiStd(diopiContextHandle_t ctx, diopiTensorHandle_t out, diopiConstTensorHandle_t input, diopiSize_t dim, bool unbiased);
+DIOPI_API diopiError_t diopiStd(diopiContextHandle_t ctx, diopiTensorHandle_t out, diopiConstTensorHandle_t input, diopiSize_t dim,
+                                const diopiScalar_t* correction);
 
 /**
  * @brief Return the minimum value of each row in the input tensor along the given dimension dim.