add pnp domain with one op: LinalgSVD #5821
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,37 @@ | ||
| # Copyright (c) ONNX Project Contributors | ||
|
Check warning Code scanning / lintrunner BLACK-ISORT/format Warning |
||
|
|
||
| # SPDX-License-Identifier: Apache-2.0 | ||
|
|
||
| import numpy as np | ||
|
|
||
| import onnx | ||
| from onnx.backend.test.case.base import Base | ||
| from onnx.backend.test.case.node import expect | ||
|
|
||
|
|
||
| class LinalgSVD(Base): | ||
| @staticmethod | ||
| def export() -> None: | ||
| threshold = 1.0 | ||
| node = onnx.helper.make_node( | ||
| "LinalgSVD", | ||
| inputs=["A"], | ||
| outputs=["U", "S", "Vh"], | ||
| threshold=threshold, | ||
|
|
||
| domain="ai.onnx.pnp", | ||
| ) | ||
|
|
||
| A = np.array([ | ||
| [ | ||
| [-1.125840, -1.152360, -0.250579, -0.433879], | ||
| [0.848710, 0.692009, -0.316013, -2.115219], | ||
| [0.468096, -0.157712, 1.443660, 0.266049], | ||
| ], | ||
| [ | ||
| [0.166455, 0.874382, -0.143474, -0.111609], | ||
| [0.931827, 1.259009, 2.004981, 0.053737], | ||
| [0.618057, -0.412802, -0.841065, -2.316042] | ||
| ] | ||
| ]) | ||
| U, S, Vh = np.linalg.svd(A, full_matrices=True) | ||
|
Check warning Code scanning / lintrunner EDITORCONFIG-CHECKER/editorconfig Warning Check warning Code scanning / lintrunner RUFF/W291 Warning |
||
| expect(node, inputs=[A], outputs=[U, S, Vh], name="test_ai_onnx_pnp_linalg_svd") | ||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,31 @@ | ||
| /* | ||
| * SPDX-License-Identifier: Apache-2.0 | ||
| */ | ||
|
|
||
| #pragma once | ||
|
|
||
| #include "onnx/defs/schema.h" | ||
|
|
||
| namespace ONNX_NAMESPACE { | ||
|
|
||
| // Declare training operators. | ||
|
|
||
| class ONNX_PNP_OPERATOR_SET_SCHEMA_CLASS_NAME(1, LinalgSVD); | ||
|
|
||
| // Iterate over schema from ai.onnx.training version 1 | ||
| class OpSet_OnnxPNP_ver1 { | ||
| public: | ||
| static void ForEachSchema(std::function<void(OpSchema&&)> fn) { | ||
|
Check warning on line 18 in onnx/defs/operator_sets_pnp.h
|
||
| fn(GetOpSchema<ONNX_PNP_OPERATOR_SET_SCHEMA_CLASS_NAME(1, LinalgSVD)>()); | ||
| } | ||
| }; | ||
|
|
||
| // Register preview operators. | ||
| inline void RegisterOnnxPNPOperatorSetSchema() { | ||
| // Preview operators should have only one version. | ||
| // If changes are needed for a specific preview operator, | ||
| // its spec should be modified without increasing its version. | ||
| RegisterOpSetSchema<OpSet_OnnxPNP_ver1>(); | ||
| } | ||
|
|
||
| } // namespace ONNX_NAMESPACE | ||
|
Check warning on line 31 in onnx/defs/operator_sets_pnp.h
|
||
GitHub Actions | Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,75 @@ | ||
| /* | ||
|
Check warning Code scanning / lintrunner CLANGFORMAT/format Warning Check warning Code scanning / lintrunner EDITORCONFIG-CHECKER/editorconfig Warning |
||
| * SPDX-License-Identifier: Apache-2.0 | ||
| */ | ||
|
|
||
| #include <algorithm> | ||
| #include <cmath> | ||
| #include <numeric> | ||
|
|
||
| #include "onnx/defs/schema.h" | ||
|
|
||
| namespace ONNX_NAMESPACE { | ||
|
|
||
| ONNX_PNP_OPERATOR_SET_SCHEMA(LinalgSVD, 1, | ||
|
Check warning on line 13 in onnx/defs/pnp/def.cc
|
||
|
There was a problem hiding this comment. We could just call it "SVD". Especially if we change the domain to "ai.onnx.linalg" as discussed elsewhere. |
||
| OpSchema() | ||
| .SetDoc(R"DOC(For internal use.)DOC") | ||
|
There was a problem hiding this comment. That would be great to have a longer description. onnx runtimes do not need to declare an operator in onnx to implement a kernel. Is it needed for onnx? There was a problem hiding this comment. +1. I don't understand what "internal use" means here. |
||
| .Attr( | ||
| "full_matrices", | ||
| "", | ||
| AttributeProto::INT, | ||
| static_cast<int64_t>(1)) | ||
| .Input( | ||
| 0, | ||
| "A", | ||
| "", | ||
| "T") | ||
| .Output( | ||
| 0, | ||
| "U", | ||
| "", | ||
| "T") | ||
| .Output( | ||
| 1, | ||
| "S", | ||
| "", | ||
| "T") | ||
| .Output( | ||
| 2, | ||
| "Vh", | ||
| "", | ||
| "T") | ||
| .TypeConstraint( | ||
| "T", | ||
| {"tensor(float)", "tensor(double)"}, | ||
| "") | ||
| .TypeAndShapeInferenceFunction([](ONNX_NAMESPACE::InferenceContext& ctx) { | ||
| ONNX_NAMESPACE::propagateElemTypeFromInputToOutput(ctx, 0, 0); | ||
| ONNX_NAMESPACE::propagateElemTypeFromInputToOutput(ctx, 0, 1); | ||
| ONNX_NAMESPACE::propagateElemTypeFromInputToOutput(ctx, 0, 2); | ||
| int64_t full_matrices = ctx.getAttribute("full_matrices")->i(); | ||
|
|
||
| const TensorShapeProto& A_shape = ctx.getInputType(0)->tensor_type().shape(); | ||
|
There was a problem hiding this comment. These lines should be protected by the usual |
||
| const auto& M = A_shape.dim(A_shape.dim_size() - 2); | ||
| const auto& N = A_shape.dim(A_shape.dim_size() - 1); | ||
| if (!M.has_dim_value() || !N.has_dim_value()) { | ||
| // cannot do shape inference without knowing dimension values | ||
| return; | ||
| } | ||
| const auto& K = M.dim_value() < N.dim_value() ? M : N; | ||
| auto* u_shape = ctx.getOutputType(0)->mutable_tensor_type()->mutable_shape(); | ||
| auto* s_shape = ctx.getOutputType(1)->mutable_tensor_type()->mutable_shape(); | ||
| auto* v_shape = ctx.getOutputType(2)->mutable_tensor_type()->mutable_shape(); | ||
| if (A_shape.dim_size() == 3) { | ||
|
There was a problem hiding this comment. These shape aspects should be explained in the operator documentation: I guess the input is required to be either 2-dim or 3-dim? |
||
| const auto& batch_dim = A_shape.dim(0); | ||
| *u_shape->add_dim() = batch_dim; | ||
| *s_shape->add_dim() = batch_dim; | ||
| *v_shape->add_dim() = batch_dim; | ||
| } | ||
| *u_shape->add_dim() = M; | ||
| *u_shape->add_dim() = full_matrices ? M : K; | ||
| *s_shape->add_dim() = K; | ||
| *v_shape->add_dim() = full_matrices ? N : K; | ||
| *v_shape->add_dim() = N; | ||
| })); | ||
|
|
||
| } // namespace ONNX_NAMESPACE | ||
|
Check warning on line 75 in onnx/defs/pnp/def.cc
|
||
Check warning
Code scanning / lintrunner
RUFF/format Warning