!41164 [assistant][ops][I40FJE] Add Eig operator 2022.8

Merge pull request !41164 from 李定维/Eig

mindspore/ccsrc/plugin/device/cpu/kernel/eigen/eig_cpu_kernel.cc

@@ -26,8 +26,7 @@ namespace mindspore {
 namespace kernel {
 namespace {
 constexpr size_t kInputsNum = 1;
-constexpr size_t kOutputsNumNV = 1;
-constexpr size_t kOutputsNumV = 2;
+constexpr size_t kOutputsNum = 2;
 }  // namespace
 
 void EigCpuKernelMod::InitMatrixInfo(const std::vector<size_t> &shape) {
@@ -51,16 +50,13 @@ void EigCpuKernelMod::InitMatrixInfo(const std::vector<size_t> &shape) {
 
 void EigCpuKernelMod::InitKernel(const CNodePtr &kernel_node) {
   kernel_name_ = common::AnfAlgo::GetCNodeName(kernel_node);
-  // If compute_v_ is true, then: w, v = Eig(a)
-  // If compute_v_ is false, then: w = Eig(a)
   if (common::AnfAlgo::HasNodeAttr(COMPUTE_V, kernel_node)) {
     compute_v_ = common::AnfAlgo::GetNodeAttr<bool>(kernel_node, COMPUTE_V);
   }
   size_t input_num = common::AnfAlgo::GetInputTensorNum(kernel_node);
   CHECK_KERNEL_INPUTS_NUM(input_num, kInputsNum, kernel_name_);
   size_t output_num = common::AnfAlgo::GetOutputTensorNum(kernel_node);
-  auto expect_output_num = compute_v_ ? kOutputsNumV : kOutputsNumNV;
-  CHECK_KERNEL_OUTPUTS_NUM(output_num, expect_output_num, kernel_name_);
+  CHECK_KERNEL_OUTPUTS_NUM(output_num, kOutputsNum, kernel_name_);
   auto input_shape = Convert2SizeTClipNeg(common::AnfAlgo::GetPrevNodeOutputInferShape(kernel_node, 0));
   InitMatrixInfo(input_shape);
 
@@ -101,16 +97,6 @@ bool EigCpuKernelMod::LaunchKernel(const std::vector<AddressPtr> &inputs, const
 }
 
 std::vector<std::pair<KernelAttr, EigCpuKernelMod::EigFunc>> EigCpuKernelMod::func_list_ = {
-  // If compute_v is false.
-  {KernelAttr().AddInputAttr(kNumberTypeFloat32).AddOutputAttr(kNumberTypeComplex64),
-   &EigCpuKernelMod::LaunchKernel<float, float_complex>},
-  {KernelAttr().AddInputAttr(kNumberTypeFloat64).AddOutputAttr(kNumberTypeComplex128),
-   &EigCpuKernelMod::LaunchKernel<double, double_complex>},
-  {KernelAttr().AddInputAttr(kNumberTypeComplex64).AddOutputAttr(kNumberTypeComplex64),
-   &EigCpuKernelMod::LaunchKernel<float_complex, float_complex>},
-  {KernelAttr().AddInputAttr(kNumberTypeComplex128).AddOutputAttr(kNumberTypeComplex128),
-   &EigCpuKernelMod::LaunchKernel<double_complex, double_complex>},
-  // If compute_v is true.
   {KernelAttr()
      .AddInputAttr(kNumberTypeFloat32)
      .AddOutputAttr(kNumberTypeComplex64)

mindspore/core/ops/core_ops.h

@@ -1226,6 +1226,7 @@ GVAR_DEF(PrimitivePtr, kPrimCholeskySolve, std::make_shared<Primitive>("Cholesky
 GVAR_DEF(PrimitivePtr, kPrimKLDivLossGrad, std::make_shared<Primitive>("KLDivLossGrad"));
 GVAR_DEF(PrimitivePtr, kPrimFFTWithSize, std::make_shared<Primitive>(kFFTWithSize));
 GVAR_DEF(PrimitivePtr, kPrimOrgqr, std::make_shared<Primitive>("Orgqr"));
+GVAR_DEF(PrimitivePtr, kPrimEig, std::make_shared<Primitive>("Eig"));
 
 // linalg
 GVAR_DEF(PrimitivePtr, kPrimGeqrf, std::make_shared<Primitive>("Geqrf"));

mindspore/core/ops/eig.cc

@@ -0,0 +1,99 @@
+/**
+ * Copyright 2021 Huawei Technologies Co., Ltd
+ *
+ * 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 "ops/eig.h"
+
+#include "abstract/ops/primitive_infer_map.h"
+#include "mindapi/src/helper.h"
+#include "ops/op_utils.h"
+#include "utils/check_convert_utils.h"
+
+namespace mindspore {
+namespace ops {
+namespace {
+abstract::TupleShapePtr EigInferShape(const PrimitivePtr &primitive, const std::vector<AbstractBasePtr> &input_args) {
+  MS_EXCEPTION_IF_NULL(primitive);
+  auto op_name = primitive->name();
+  auto input_x = CheckAndConvertUtils::CheckArgs<abstract::AbstractTensor>(op_name, input_args, kInputIndex0);
+  auto x_shape = input_x->shape();
+  MS_EXCEPTION_IF_NULL(x_shape);
+  constexpr size_t kDefaultRank = 2;
+  constexpr size_t kRowIndex = 2;
+  constexpr size_t kColIndex = 1;
+  auto const &x_shape_list = x_shape->shape();
+  const size_t x_rank = x_shape_list.size();
+  if (x_rank < kDefaultRank) {
+    MS_EXCEPTION(ValueError) << "For Eig, x should be at least rank 2"
+                             << ", but got a " << x_rank << "-D Tensor.";
+  }
+  if (x_shape_list[x_rank - kRowIndex] != x_shape_list[x_rank - kColIndex]) {
+    MS_EXCEPTION(ValueError) << "For Eig, x should be square(squares)"
+                             << ", but got " << x_shape_list[x_rank - kRowIndex] << " × "
+                             << x_shape_list[x_rank - kColIndex] << " matrix(matrices).";
+  }
+  auto compute_v = GetValue<bool>(primitive->GetAttr("compute_v"));
+  std::vector<BaseShapePtr> shapes_list;
+  if (compute_v) {
+    ShapeVector val_shape_list;
+    val_shape_list.assign(x_shape_list.begin(), x_shape_list.end());
+    val_shape_list.pop_back();
+    (void)shapes_list.emplace_back(std::make_shared<abstract::Shape>(val_shape_list));
+    (void)shapes_list.emplace_back(std::make_shared<abstract::Shape>(x_shape_list));
+    return std::make_shared<abstract::TupleShape>(shapes_list);
+  } else {
+    ShapeVector val_shape_list;
+    val_shape_list.assign(x_shape_list.begin(), x_shape_list.end());
+    val_shape_list.pop_back();
+    ShapeVector empyty_shape_list = {};
+    (void)shapes_list.emplace_back(std::make_shared<abstract::Shape>(val_shape_list));
+    (void)shapes_list.emplace_back(std::make_shared<abstract::Shape>(empyty_shape_list));
+    return std::make_shared<abstract::TupleShape>(shapes_list);
+  }
+}
+
+TuplePtr EigInferType(const PrimitivePtr &primitive, const std::vector<AbstractBasePtr> &input_args) {
+  MS_EXCEPTION_IF_NULL(primitive);
+  auto op_name = primitive->name();
+  const std::set<TypePtr> valid_types = {kFloat32, kFloat64, kComplex64, kComplex128};
+  auto x_type = input_args[kInputIndex0]->BuildType();
+  (void)CheckAndConvertUtils::CheckTensorTypeValid("x", x_type, valid_types, op_name);
+  std::vector<TypePtr> types_list;
+  if (*(x_type->cast<TensorTypePtr>()->element()) == *(kFloat32)) {
+    types_list = {kComplex64, kComplex64};
+  } else if (*(x_type->cast<TensorTypePtr>()->element()) == *(kFloat64)) {
+    types_list = {kComplex128, kComplex128};
+  } else {
+    types_list = {x_type, x_type};
+  }
+
+  return std::make_shared<Tuple>(types_list);
+}
+}  // namespace
+
+MIND_API_OPERATOR_IMPL(Eig, BaseOperator);
+AbstractBasePtr EigInfer(const abstract::AnalysisEnginePtr &, const PrimitivePtr &primitive,
+                         const std::vector<AbstractBasePtr> &input_args) {
+  MS_EXCEPTION_IF_NULL(primitive);
+  auto op_name = primitive->name();
+  const int64_t input_num = 1;
+  (void)CheckAndConvertUtils::CheckInputArgs(input_args, kEqual, input_num, op_name);
+  auto infer_type = EigInferType(primitive, input_args);
+  auto infer_shape = EigInferShape(primitive, input_args);
+  return abstract::MakeAbstract(infer_shape, infer_type);
+}
+REGISTER_PRIMITIVE_EVAL_IMPL(Eig, prim::kPrimEig, EigInfer, nullptr, true);
+}  // namespace ops
+}  // namespace mindspore

mindspore/core/ops/eig.h

@@ -0,0 +1,45 @@
+/**
+ * Copyright 2021 Huawei Technologies Co., Ltd
+ *
+ * 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.
+ */
+
+#ifndef MINDSPORE_CORE_OPS_EIG_H_
+#define MINDSPORE_CORE_OPS_EIG_H_
+
+#include <algorithm>
+#include <map>
+#include <memory>
+#include <set>
+#include <string>
+#include <vector>
+#include "ops/base_operator.h"
+#include "mindapi/base/types.h"
+
+namespace mindspore {
+namespace ops {
+/// \brief Computes the eigenvalue decomposition of a (batched) square matrix.
+/// Refer to Python API @ref mindspore.ops.Eig for more details.
+constexpr auto kNameEig = "Eig";
+class MIND_API Eig : public BaseOperator {
+ public:
+  MIND_API_BASE_MEMBER(Eig);
+  /// \brief Constructor.
+  Eig() : BaseOperator(kNameEig) { InitIOName({"x"}, {"eigen_values", "eigen_vectors"}); }
+};
+abstract::AbstractBasePtr EigInfer(const abstract::AnalysisEnginePtr &, const PrimitivePtr &primitive,
+                                   const std::vector<abstract::AbstractBasePtr> &input_args);
+}  // namespace ops
+}  // namespace mindspore
+
+#endif  // MINDSPORE_CORE_OPS_EIG_H_

mindspore/python/mindspore/ops/_op_impl/aicpu/eig.py

@@ -0,0 +1,35 @@
+# Copyright 2021 Huawei Technologies Co., Ltd
+#
+# 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.
+# ============================================================================
+
+"""Eig op"""
+from mindspore.ops.op_info_register import op_info_register, AiCPURegOp, DataType
+
+eig_op_info = AiCPURegOp("Eig") \
+    .fusion_type("OPAQUE") \
+    .attr("compute_v", "bool") \
+    .input(0, "x", "required") \
+    .output(0, "eigen_values", "required") \
+    .output(1, "eigen_vectors", "required") \
+    .dtype_format(DataType.F32_Default, DataType.C64_Default, DataType.C64_Default) \
+    .dtype_format(DataType.C64_Default, DataType.C64_Default, DataType.C64_Default) \
+    .dtype_format(DataType.F64_Default, DataType.C128_Default, DataType.C128_Default) \
+    .dtype_format(DataType.C128_Default, DataType.C128_Default, DataType.C128_Default) \
+    .get_op_info()
+
+
+@op_info_register(eig_op_info)
+def _eig_aicpu():
+    """Eig AiCPU register"""
+    return

mindspore/python/mindspore/ops/operations/math_ops.py

@@ -6857,3 +6857,48 @@ class Orgqr(Primitive):
     def __init__(self):
         """Initialize Orgqr"""
         self.init_prim_io_names(inputs=['x', 'tau'], outputs=['y'])
+
+
+class Eig(Primitive):
+    """
+    Computes the eigenvalues and eigenvectors of a square matrix(batch square matrices).
+
+    Args:
+        compute_v (bool): If `True`, compute both eigenvalues and eigenvectors;
+            If `False`, just eigenvalues will be computed. Default: False.
+    Inputs:
+        - **x** (Tensor) - Square matrices of shape :math:`(*, N, N)`,
+          with float32, float64, complex64 or complex128 data type.
+
+    Outputs:
+        - **eigen_values** (Tensor) - Shape :math:`(*, N)`. Each inner most vector represents eigenvalues of
+          the corresponding matrix. The eigenvalues may not have an order.
+        - **eigen_vectors** (Tensor) - If `compute_v` is `False`, it’s an empty tensor. Otherwise, this tensor
+          has shape :math:`(*, N, N)`, whose columns represent normalized (unit length) eigenvectors of corresponding
+          eigenvalues.
+
+    Raises:
+        TypeError: If `compute_v` is not a bool.
+        TypeError: If dtype of `x` is not one of: float64, float32, complex64 or complex128.
+        TypeError: If `x` is not a Tensor.
+        ValueError: If `x` is not a square(batch squares).
+
+    Supported Platforms:
+        ``Ascend`` ``CPU``
+
+    Examples:
+        >>> input_x = Tensor(np.array([[1.0, 0.0], [0.0, 2.0]]), mindspore.float32)
+        >>> eig = ops.Eig(compute_v=True)
+        >>> u, v = eig(input_x)
+        >>> print(u)
+        [1.+0.j 2.+0.j]
+        >>> print(v)
+        [[1.+0.j 0.+0.j]
+         [0.+0.j 1.+0.j]]
+    """
+
+    @prim_attr_register
+    def __init__(self, compute_v=False):
+        """Initialize Eig"""
+        self.init_prim_io_names(inputs=['x'], outputs=['eigen_values', 'eigen_vectors'])
+        validator.check_value_type('compute_v', compute_v, [bool], self.name)

tests/st/scipy_st/test_ops.py

@@ -112,7 +112,7 @@ def test_eig(shape, data_type, rtol, atol):
     compare_eigen_decomposition((mw, mv), (sw, sv), True, rtol, atol)
 
     # Eig only calculate eigenvalues when compute_v is False
-    mw = Eig(False)(tensor_a)
+    mw, _ = Eig(False)(tensor_a)
     mw = mw.asnumpy()
     sw = eigvals(a)
     compare_eigen_decomposition((mw,), (sw,), False, rtol, atol)

tests/ut/python/ops/test_ops.py

@@ -42,6 +42,7 @@ from mindspore.ops.operations.array_ops import ConjugateTranspose
 from mindspore.ops.operations.array_ops import UnravelIndex
 from mindspore.ops.operations.math_ops import Trace
 from mindspore.ops.operations.math_ops import Cholesky
+from mindspore.ops.operations.math_ops import Eig
 from mindspore.ops.operations.math_ops import LuUnpack
 from mindspore.ops.operations.math_ops import MatrixExp
 from mindspore.ops.operations.math_ops import MatrixSolve
@@ -1453,6 +1454,10 @@ test_case_math_ops = [
         'desc_inputs': [Tensor(np.array([0, 0, 1, -1, 1, 1, 1]), mstype.int16),
                         Tensor(np.array([0, 1, 1, -1, -1, 2, 3]), mstype.int16)],
         'skip': ['backward']}),
+    ('Eig', {
+        'block': Eig(),
+        'desc_inputs': [Tensor(np.array([[1, 0], [0, 1]]).astype(np.float32))],
+        'skip': ['backward']}),
     ('BitwiseOr_1', {
         'block': P.BitwiseOr(),
         'desc_inputs': [Tensor(np.array([[1, 2, 3], [-1, -2, -3]]), mstype.int16),