| /* |
| * Licensed to the Apache Software Foundation (ASF) under one |
| * or more contributor license agreements. See the NOTICE file |
| * distributed with this work for additional information |
| * regarding copyright ownership. The ASF licenses this file |
| * to you 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. |
| */ |
| |
| /*! |
| * \file la_op.cc |
| * \brief CPU implementation of Operators for advanced linear algebra. |
| */ |
| |
| #include "./la_op.h" |
| #include "./la_op-inl.h" |
| |
| namespace mxnet { |
| namespace op { |
| |
| DMLC_REGISTER_PARAMETER(LaMatrixMacParam); |
| DMLC_REGISTER_PARAMETER(LaMatrixMultParam); |
| DMLC_REGISTER_PARAMETER(LaCholeskyParam); |
| DMLC_REGISTER_PARAMETER(LaTriangMatrixMultParam); |
| DMLC_REGISTER_PARAMETER(LaDiagParam); |
| DMLC_REGISTER_PARAMETER(LaTrianParam); |
| DMLC_REGISTER_PARAMETER(LaSyrkParam); |
| |
| NNVM_REGISTER_OP(_linalg_gemm) |
| .add_alias("linalg_gemm") |
| .describe(R"code(Performs general matrix multiplication and accumulation. |
| Input are tensors *A*, *B*, *C*, each of dimension *n >= 2* and having the same shape |
| on the leading *n-2* dimensions. |
| |
| If *n=2*, the BLAS3 function *gemm* is performed: |
| |
| *out* = *alpha* \* *op*\ (*A*) \* *op*\ (*B*) + *beta* \* *C* |
| |
| Here, *alpha* and *beta* are scalar parameters, and *op()* is either the identity or |
| matrix transposition (depending on *transpose_a*, *transpose_b*). |
| |
| If *n>2*, *gemm* is performed separately for a batch of matrices. The column indices of the matrices |
| are given by the last dimensions of the tensors, the row indices by the axis specified with the *axis* |
| parameter. By default, the trailing two dimensions will be used for matrix encoding. |
| |
| For a non-default axis parameter, the operation performed is equivalent to a series of swapaxes/gemm/swapaxes |
| calls. For example let *A*, *B*, *C* be 5 dimensional tensors. Then gemm(*A*, *B*, *C*, axis=1) is equivalent |
| to the following without the overhead of the additional swapaxis operations:: |
| |
| A1 = swapaxes(A, dim1=1, dim2=3) |
| B1 = swapaxes(B, dim1=1, dim2=3) |
| C = swapaxes(C, dim1=1, dim2=3) |
| C = gemm(A1, B1, C) |
| C = swapaxis(C, dim1=1, dim2=3) |
| |
| When the input data is of type float32 and the environment variables MXNET_CUDA_ALLOW_TENSOR_CORE |
| and MXNET_CUDA_TENSOR_OP_MATH_ALLOW_CONVERSION are set to 1, this operator will try to use |
| pseudo-float16 precision (float32 math with float16 I/O) precision in order to use |
| Tensor Cores on suitable NVIDIA GPUs. This can sometimes give significant speedups. |
| |
| .. note:: The operator supports float32 and float64 data types only. |
| |
| Examples:: |
| |
| Single matrix multiply-add |
| A = [[1.0, 1.0], [1.0, 1.0]] |
| B = [[1.0, 1.0], [1.0, 1.0], [1.0, 1.0]] |
| C = [[1.0, 1.0, 1.0], [1.0, 1.0, 1.0]] |
| gemm(A, B, C, transpose_b=True, alpha=2.0, beta=10.0) |
| = [[14.0, 14.0, 14.0], [14.0, 14.0, 14.0]] |
| |
| Batch matrix multiply-add |
| A = [[[1.0, 1.0]], [[0.1, 0.1]]] |
| B = [[[1.0, 1.0]], [[0.1, 0.1]]] |
| C = [[[10.0]], [[0.01]]] |
| gemm(A, B, C, transpose_b=True, alpha=2.0 , beta=10.0) |
| = [[[104.0]], [[0.14]]] |
| )code" ADD_FILELINE) |
| .set_num_inputs(3) |
| .set_num_outputs(1) |
| .set_attr_parser(ParamParser<LaMatrixMacParam>) |
| .set_attr<nnvm::FListInputNames>("FListInputNames", |
| [](const NodeAttrs& attrs) { |
| return std::vector<std::string>{"A", "B", "C"}; |
| }) |
| .set_attr<mxnet::FInferShape>("FInferShape", LaMatrixMultMacOpShape) |
| .set_attr<nnvm::FInferType>("FInferType", ElemwiseType<3, 1>) |
| .set_attr<nnvm::FInplaceOption>("FInplaceOption", |
| [](const NodeAttrs& attrs) { |
| return std::vector<std::pair<int, int>>{{2, 0}}; |
| }) |
| .set_attr<FCompute>("FCompute<cpu>", LaOpGemmForward<cpu, 2, 2, 3, 1, gemm>) |
| .set_attr<nnvm::FGradient>("FGradient", ElemwiseGradUseIn{"_backward_linalg_gemm"}) |
| .add_argument("A", "NDArray-or-Symbol", "Tensor of input matrices") |
| .add_argument("B", "NDArray-or-Symbol", "Tensor of input matrices") |
| .add_argument("C", "NDArray-or-Symbol", "Tensor of input matrices") |
| .add_arguments(LaMatrixMacParam::__FIELDS__()); |
| |
| NNVM_REGISTER_OP(_backward_linalg_gemm) |
| .set_num_inputs(4) |
| .set_num_outputs(3) |
| .set_attr_parser(ParamParser<LaMatrixMacParam>) |
| .set_attr<nnvm::FInplaceOption>("FInplaceOption", |
| [](const NodeAttrs& attrs) { |
| return std::vector<std::pair<int, int>>{{2, 1}, {3, 2}}; |
| }) |
| .set_attr<FResourceRequest>("FResourceRequest", |
| [](const NodeAttrs& attrs) { |
| return std::vector<ResourceRequest>{ResourceRequest::kTempSpace}; |
| }) |
| .set_attr<nnvm::TIsBackward>("TIsBackward", true) |
| .set_attr<FCompute>("FCompute<cpu>", LaOpGemmBackward<cpu, 2, 2, 4, 3, gemm_backward>); |
| |
| NNVM_REGISTER_OP(_linalg_gemm2) |
| .add_alias("linalg_gemm2") |
| .describe(R"code(Performs general matrix multiplication. |
| Input are tensors *A*, *B*, each of dimension *n >= 2* and having the same shape |
| on the leading *n-2* dimensions. |
| |
| If *n=2*, the BLAS3 function *gemm* is performed: |
| |
| *out* = *alpha* \* *op*\ (*A*) \* *op*\ (*B*) |
| |
| Here *alpha* is a scalar parameter and *op()* is either the identity or the matrix |
| transposition (depending on *transpose_a*, *transpose_b*). |
| |
| If *n>2*, *gemm* is performed separately for a batch of matrices. The column indices of the matrices |
| are given by the last dimensions of the tensors, the row indices by the axis specified with the *axis* |
| parameter. By default, the trailing two dimensions will be used for matrix encoding. |
| |
| For a non-default axis parameter, the operation performed is equivalent to a series of swapaxes/gemm/swapaxes |
| calls. For example let *A*, *B* be 5 dimensional tensors. Then gemm(*A*, *B*, axis=1) is equivalent to |
| the following without the overhead of the additional swapaxis operations:: |
| |
| A1 = swapaxes(A, dim1=1, dim2=3) |
| B1 = swapaxes(B, dim1=1, dim2=3) |
| C = gemm2(A1, B1) |
| C = swapaxis(C, dim1=1, dim2=3) |
| |
| When the input data is of type float32 and the environment variables MXNET_CUDA_ALLOW_TENSOR_CORE |
| and MXNET_CUDA_TENSOR_OP_MATH_ALLOW_CONVERSION are set to 1, this operator will try to use |
| pseudo-float16 precision (float32 math with float16 I/O) precision in order to use |
| Tensor Cores on suitable NVIDIA GPUs. This can sometimes give significant speedups. |
| |
| .. note:: The operator supports float32 and float64 data types only. |
| |
| Examples:: |
| |
| Single matrix multiply |
| A = [[1.0, 1.0], [1.0, 1.0]] |
| B = [[1.0, 1.0], [1.0, 1.0], [1.0, 1.0]] |
| gemm2(A, B, transpose_b=True, alpha=2.0) |
| = [[4.0, 4.0, 4.0], [4.0, 4.0, 4.0]] |
| |
| Batch matrix multiply |
| A = [[[1.0, 1.0]], [[0.1, 0.1]]] |
| B = [[[1.0, 1.0]], [[0.1, 0.1]]] |
| gemm2(A, B, transpose_b=True, alpha=2.0) |
| = [[[4.0]], [[0.04 ]]] |
| )code" ADD_FILELINE) |
| .set_num_inputs(2) |
| .set_num_outputs(1) |
| .set_attr_parser(ParamParser<LaMatrixMultParam>) |
| .set_attr<nnvm::FListInputNames>("FListInputNames", |
| [](const NodeAttrs& attrs) { |
| return std::vector<std::string>{"A", "B"}; |
| }) |
| .set_attr<mxnet::FInferShape>("FInferShape", LaMatrixMultMacOpShape) |
| .set_attr<nnvm::FInferType>("FInferType", ElemwiseType<2, 1>) |
| .set_attr<FCompute>("FCompute<cpu>", LaOpGemmForward<cpu, 2, 2, 2, 1, gemm2>) |
| .set_attr<nnvm::FGradient>("FGradient", ElemwiseGradUseIn{"_backward_linalg_gemm2"}) |
| .add_argument("A", "NDArray-or-Symbol", "Tensor of input matrices") |
| .add_argument("B", "NDArray-or-Symbol", "Tensor of input matrices") |
| .add_arguments(LaMatrixMultParam::__FIELDS__()); |
| |
| NNVM_REGISTER_OP(_backward_linalg_gemm2) |
| .set_num_inputs(3) |
| .set_num_outputs(2) |
| .set_attr_parser(ParamParser<LaMatrixMultParam>) |
| .set_attr<nnvm::FInplaceOption>("FInplaceOption", |
| [](const NodeAttrs& attrs) { |
| return std::vector<std::pair<int, int>>{{2, 1}}; |
| }) |
| .set_attr<FResourceRequest>("FResourceRequest", |
| [](const NodeAttrs& attrs) { |
| return std::vector<ResourceRequest>{ResourceRequest::kTempSpace}; |
| }) |
| .set_attr<nnvm::TIsBackward>("TIsBackward", true) |
| .set_attr<FCompute>("FCompute<cpu>", LaOpGemmBackward<cpu, 2, 2, 3, 2, gemm2_backward>); |
| |
| NNVM_REGISTER_OP(_linalg_potrf) |
| .add_alias("linalg_potrf") |
| .describe(R"code(Performs Cholesky factorization of a symmetric positive-definite matrix. |
| Input is a tensor *A* of dimension *n >= 2*. |
| |
| If *n=2*, the Cholesky factor *B* of the symmetric, positive definite matrix *A* is |
| computed. *B* is triangular (entries of upper or lower triangle are all zero), has |
| positive diagonal entries, and: |
| |
| *A* = *B* \* *B*\ :sup:`T` if *lower* = *true* |
| *A* = *B*\ :sup:`T` \* *B* if *lower* = *false* |
| |
| If *n>2*, *potrf* is performed separately on the trailing two dimensions for all inputs |
| (batch mode). |
| |
| .. note:: The operator supports float32 and float64 data types only. |
| |
| Examples:: |
| |
| Single matrix factorization |
| A = [[4.0, 1.0], [1.0, 4.25]] |
| potrf(A) = [[2.0, 0], [0.5, 2.0]] |
| |
| Batch matrix factorization |
| A = [[[4.0, 1.0], [1.0, 4.25]], [[16.0, 4.0], [4.0, 17.0]]] |
| potrf(A) = [[[2.0, 0], [0.5, 2.0]], [[4.0, 0], [1.0, 4.0]]] |
| )code" ADD_FILELINE) |
| .set_num_inputs(1) |
| .set_num_outputs(1) |
| .set_attr_parser(ParamParser<LaCholeskyParam>) |
| .set_attr<nnvm::FListInputNames>("FListInputNames", |
| [](const NodeAttrs& attrs) { |
| return std::vector<std::string>{"A"}; |
| }) |
| .set_attr<mxnet::FInferShape>("FInferShape", ElemwiseShape<1, 1>) |
| .set_attr<nnvm::FInferType>("FInferType", ElemwiseType<1, 1>) |
| .set_attr<nnvm::FInplaceOption>("FInplaceOption", |
| [](const NodeAttrs& attrs) { |
| return std::vector<std::pair<int, int>>{{0, 0}}; |
| }) |
| .set_attr<FCompute>("FCompute<cpu>", LaOpForward<cpu, 2, 2, 1, 1, potrf>) |
| .set_attr<nnvm::FGradient>("FGradient", ElemwiseGradUseOut{"_backward_linalg_potrf"}) |
| .add_argument("A", "NDArray-or-Symbol", "Tensor of input matrices to be decomposed"); |
| |
| NNVM_REGISTER_OP(_backward_linalg_potrf) |
| .set_num_inputs(2) |
| .set_num_outputs(1) |
| .set_attr_parser(ParamParser<LaCholeskyParam>) |
| .set_attr<nnvm::FInplaceOption>("FInplaceOption", |
| [](const NodeAttrs& attrs) { |
| return std::vector<std::pair<int, int>>{{0, 0}}; |
| }) |
| .set_attr<FResourceRequest>("FResourceRequest", |
| [](const NodeAttrs& attrs) { |
| return std::vector<ResourceRequest>{ResourceRequest::kTempSpace}; |
| }) |
| .set_attr<nnvm::TIsBackward>("TIsBackward", true) |
| .set_attr<FCompute>("FCompute<cpu>", LaOpBackward<cpu, 2, 2, 2, 1, potrf_backward>); |
| |
| NNVM_REGISTER_OP(_linalg_potri) |
| .add_alias("linalg_potri") |
| .describe(R"code(Performs matrix inversion from a Cholesky factorization. |
| Input is a tensor *A* of dimension *n >= 2*. |
| |
| If *n=2*, *A* is a triangular matrix (entries of upper or lower triangle are all zero) |
| with positive diagonal. We compute: |
| |
| *out* = *A*\ :sup:`-T` \* *A*\ :sup:`-1` if *lower* = *true* |
| *out* = *A*\ :sup:`-1` \* *A*\ :sup:`-T` if *lower* = *false* |
| |
| In other words, if *A* is the Cholesky factor of a symmetric positive definite matrix |
| *B* (obtained by *potrf*), then |
| |
| *out* = *B*\ :sup:`-1` |
| |
| If *n>2*, *potri* is performed separately on the trailing two dimensions for all inputs |
| (batch mode). |
| |
| .. note:: The operator supports float32 and float64 data types only. |
| |
| .. note:: Use this operator only if you are certain you need the inverse of *B*, and |
| cannot use the Cholesky factor *A* (*potrf*), together with backsubstitution |
| (*trsm*). The latter is numerically much safer, and also cheaper. |
| |
| Examples:: |
| |
| Single matrix inverse |
| A = [[2.0, 0], [0.5, 2.0]] |
| potri(A) = [[0.26563, -0.0625], [-0.0625, 0.25]] |
| |
| Batch matrix inverse |
| A = [[[2.0, 0], [0.5, 2.0]], [[4.0, 0], [1.0, 4.0]]] |
| potri(A) = [[[0.26563, -0.0625], [-0.0625, 0.25]], |
| [[0.06641, -0.01562], [-0.01562, 0,0625]]] |
| )code" ADD_FILELINE) |
| .set_num_inputs(1) |
| .set_num_outputs(1) |
| .set_attr_parser(ParamParser<LaCholeskyParam>) |
| .set_attr<nnvm::FListInputNames>("FListInputNames", |
| [](const NodeAttrs& attrs) { |
| return std::vector<std::string>{"A"}; |
| }) |
| .set_attr<mxnet::FInferShape>("FInferShape", ElemwiseShape<1, 1>) |
| .set_attr<nnvm::FInferType>("FInferType", ElemwiseType<1, 1>) |
| .set_attr<nnvm::FInplaceOption>("FInplaceOption", |
| [](const NodeAttrs& attrs) { |
| return std::vector<std::pair<int, int>>{{0, 0}}; |
| }) |
| .set_attr<FCompute>("FCompute<cpu>", LaOpForward<cpu, 2, 2, 1, 1, potri>) |
| .set_attr<nnvm::FGradient>("FGradient", ElemwiseGradUseInOut{"_backward_linalg_potri"}) |
| .add_argument("A", "NDArray-or-Symbol", "Tensor of lower triangular matrices"); |
| |
| NNVM_REGISTER_OP(_backward_linalg_potri) |
| .set_num_inputs(3) |
| .set_num_outputs(1) |
| .set_attr_parser(ParamParser<LaCholeskyParam>) |
| .set_attr<FResourceRequest>("FResourceRequest", |
| [](const NodeAttrs& attrs) { |
| return std::vector<ResourceRequest>{ResourceRequest::kTempSpace}; |
| }) |
| .set_attr<nnvm::TIsBackward>("TIsBackward", true) |
| .set_attr<FCompute>("FCompute<cpu>", LaOpBackward<cpu, 2, 2, 3, 1, potri_backward>); |
| |
| NNVM_REGISTER_OP(_linalg_trmm) |
| .add_alias("linalg_trmm") |
| .describe(R"code(Performs multiplication with a lower triangular matrix. |
| Input are tensors *A*, *B*, each of dimension *n >= 2* and having the same shape |
| on the leading *n-2* dimensions. |
| |
| If *n=2*, *A* must be triangular. The operator performs the BLAS3 function |
| *trmm*: |
| |
| *out* = *alpha* \* *op*\ (*A*) \* *B* |
| |
| if *rightside=False*, or |
| |
| *out* = *alpha* \* *B* \* *op*\ (*A*) |
| |
| if *rightside=True*. Here, *alpha* is a scalar parameter, and *op()* is either the |
| identity or the matrix transposition (depending on *transpose*). |
| |
| If *n>2*, *trmm* is performed separately on the trailing two dimensions for all inputs |
| (batch mode). |
| |
| .. note:: The operator supports float32 and float64 data types only. |
| |
| Examples:: |
| |
| Single triangular matrix multiply |
| A = [[1.0, 0], [1.0, 1.0]] |
| B = [[1.0, 1.0, 1.0], [1.0, 1.0, 1.0]] |
| trmm(A, B, alpha=2.0) = [[2.0, 2.0, 2.0], [4.0, 4.0, 4.0]] |
| |
| Batch triangular matrix multiply |
| A = [[[1.0, 0], [1.0, 1.0]], [[1.0, 0], [1.0, 1.0]]] |
| B = [[[1.0, 1.0, 1.0], [1.0, 1.0, 1.0]], [[0.5, 0.5, 0.5], [0.5, 0.5, 0.5]]] |
| trmm(A, B, alpha=2.0) = [[[2.0, 2.0, 2.0], [4.0, 4.0, 4.0]], |
| [[1.0, 1.0, 1.0], [2.0, 2.0, 2.0]]] |
| )code" ADD_FILELINE) |
| .set_num_inputs(2) |
| .set_num_outputs(1) |
| .set_attr_parser(ParamParser<LaTriangMatrixMultParam>) |
| .set_attr<nnvm::FListInputNames>("FListInputNames", |
| [](const NodeAttrs& attrs) { |
| return std::vector<std::string>{"A", "B"}; |
| }) |
| .set_attr<mxnet::FInferShape>("FInferShape", LaTriangMatrixMultOpShape) |
| .set_attr<nnvm::FInferType>("FInferType", ElemwiseType<2, 1>) |
| .set_attr<nnvm::FInplaceOption>("FInplaceOption", |
| [](const NodeAttrs& attrs) { |
| return std::vector<std::pair<int, int>>{{1, 0}}; |
| }) |
| .set_attr<FCompute>("FCompute<cpu>", LaOpForward<cpu, 2, 2, 2, 1, trmm>) |
| .set_attr<nnvm::FGradient>("FGradient", ElemwiseGradUseIn{"_backward_linalg_trmm"}) |
| .add_argument("A", "NDArray-or-Symbol", "Tensor of lower triangular matrices") |
| .add_argument("B", "NDArray-or-Symbol", "Tensor of matrices") |
| .add_arguments(LaTriangMatrixMultParam::__FIELDS__()); |
| |
| NNVM_REGISTER_OP(_backward_linalg_trmm) |
| .set_num_inputs(3) |
| .set_num_outputs(2) |
| .set_attr_parser(ParamParser<LaTriangMatrixMultParam>) |
| .set_attr<nnvm::FInplaceOption>("FInplaceOption", |
| [](const NodeAttrs& attrs) { |
| return std::vector<std::pair<int, int>>{{0, 1}}; |
| }) |
| .set_attr<FResourceRequest>("FResourceRequest", |
| [](const NodeAttrs& attrs) { |
| return std::vector<ResourceRequest>{ResourceRequest::kTempSpace}; |
| }) |
| .set_attr<nnvm::TIsBackward>("TIsBackward", true) |
| .set_attr<FCompute>("FCompute<cpu>", LaOpBackward<cpu, 2, 2, 3, 2, trmm_backward>); |
| |
| NNVM_REGISTER_OP(_linalg_trsm) |
| .add_alias("linalg_trsm") |
| .describe(R"code(Solves matrix equation involving a lower triangular matrix. |
| Input are tensors *A*, *B*, each of dimension *n >= 2* and having the same shape |
| on the leading *n-2* dimensions. |
| |
| If *n=2*, *A* must be triangular. The operator performs the BLAS3 function |
| *trsm*, solving for *out* in: |
| |
| *op*\ (*A*) \* *out* = *alpha* \* *B* |
| |
| if *rightside=False*, or |
| |
| *out* \* *op*\ (*A*) = *alpha* \* *B* |
| |
| if *rightside=True*. Here, *alpha* is a scalar parameter, and *op()* is either the |
| identity or the matrix transposition (depending on *transpose*). |
| |
| If *n>2*, *trsm* is performed separately on the trailing two dimensions for all inputs |
| (batch mode). |
| |
| .. note:: The operator supports float32 and float64 data types only. |
| |
| Examples:: |
| |
| Single matrix solve |
| A = [[1.0, 0], [1.0, 1.0]] |
| B = [[2.0, 2.0, 2.0], [4.0, 4.0, 4.0]] |
| trsm(A, B, alpha=0.5) = [[1.0, 1.0, 1.0], [1.0, 1.0, 1.0]] |
| |
| Batch matrix solve |
| A = [[[1.0, 0], [1.0, 1.0]], [[1.0, 0], [1.0, 1.0]]] |
| B = [[[2.0, 2.0, 2.0], [4.0, 4.0, 4.0]], |
| [[4.0, 4.0, 4.0], [8.0, 8.0, 8.0]]] |
| trsm(A, B, alpha=0.5) = [[[1.0, 1.0, 1.0], [1.0, 1.0, 1.0]], |
| [[2.0, 2.0, 2.0], [2.0, 2.0, 2.0]]] |
| )code" ADD_FILELINE) |
| .set_num_inputs(2) |
| .set_num_outputs(1) |
| .set_attr_parser(ParamParser<LaTriangMatrixMultParam>) |
| .set_attr<nnvm::FListInputNames>("FListInputNames", |
| [](const NodeAttrs& attrs) { |
| return std::vector<std::string>{"A", "B"}; |
| }) |
| .set_attr<mxnet::FInferShape>("FInferShape", LaTriangMatrixMultOpShape) |
| .set_attr<nnvm::FInferType>("FInferType", ElemwiseType<2, 1>) |
| .set_attr<nnvm::FInplaceOption>("FInplaceOption", |
| [](const NodeAttrs& attrs) { |
| return std::vector<std::pair<int, int>>{{1, 0}}; |
| }) |
| .set_attr<FCompute>("FCompute<cpu>", LaOpForward<cpu, 2, 2, 2, 1, trsm>) |
| .set_attr<nnvm::FGradient>("FGradient", ElemwiseGradUseInOut{"_backward_linalg_trsm"}) |
| .add_argument("A", "NDArray-or-Symbol", "Tensor of lower triangular matrices") |
| .add_argument("B", "NDArray-or-Symbol", "Tensor of matrices") |
| .add_arguments(LaTriangMatrixMultParam::__FIELDS__()); |
| |
| NNVM_REGISTER_OP(_backward_linalg_trsm) |
| .set_num_inputs(4) |
| .set_num_outputs(2) |
| .set_attr_parser(ParamParser<LaTriangMatrixMultParam>) |
| .set_attr<nnvm::FInplaceOption>("FInplaceOption", |
| [](const NodeAttrs& attrs) { |
| return std::vector<std::pair<int, int>>{{0, 1}, {1, 0}}; |
| }) |
| .set_attr<FResourceRequest>("FResourceRequest", |
| [](const NodeAttrs& attrs) { |
| return std::vector<ResourceRequest>{ResourceRequest::kTempSpace}; |
| }) |
| .set_attr<nnvm::TIsBackward>("TIsBackward", true) |
| .set_attr<FCompute>("FCompute<cpu>", LaOpBackward<cpu, 2, 2, 4, 2, trsm_backward>); |
| |
| NNVM_REGISTER_OP(_linalg_sumlogdiag) |
| .add_alias("linalg_sumlogdiag") |
| .describe(R"code(Computes the sum of the logarithms of the diagonal elements of a square matrix. |
| Input is a tensor *A* of dimension *n >= 2*. |
| |
| If *n=2*, *A* must be square with positive diagonal entries. We sum the natural |
| logarithms of the diagonal elements, the result has shape (1,). |
| |
| If *n>2*, *sumlogdiag* is performed separately on the trailing two dimensions for all |
| inputs (batch mode). |
| |
| .. note:: The operator supports float32 and float64 data types only. |
| |
| Examples:: |
| |
| Single matrix reduction |
| A = [[1.0, 1.0], [1.0, 7.0]] |
| sumlogdiag(A) = [1.9459] |
| |
| Batch matrix reduction |
| A = [[[1.0, 1.0], [1.0, 7.0]], [[3.0, 0], [0, 17.0]]] |
| sumlogdiag(A) = [1.9459, 3.9318] |
| )code" ADD_FILELINE) |
| .set_num_inputs(1) |
| .set_num_outputs(1) |
| .set_attr<nnvm::FListInputNames>("FListInputNames", |
| [](const NodeAttrs& attrs) { |
| return std::vector<std::string>{"A"}; |
| }) |
| .set_attr<mxnet::FInferShape>("FInferShape", LaReduceShape<2>) |
| .set_attr<nnvm::FInferType>("FInferType", ElemwiseType<1, 1>) |
| .set_attr<FCompute>("FCompute<cpu>", LaOpForward<cpu, 2, 0, 1, 1, sumlogdiag>) |
| .set_attr<nnvm::FGradient>("FGradient", ElemwiseGradUseIn{"_backward_linalg_sumlogdiag"}) |
| .add_argument("A", "NDArray-or-Symbol", "Tensor of square matrices"); |
| |
| NNVM_REGISTER_OP(_backward_linalg_sumlogdiag) |
| .set_num_inputs(2) |
| .set_num_outputs(1) |
| .set_attr<nnvm::FInplaceOption>("FInplaceOption", |
| [](const NodeAttrs& attrs) { |
| return std::vector<std::pair<int, int>>{{1, 0}}; |
| }) |
| .set_attr<FResourceRequest>("FResourceRequest", |
| [](const NodeAttrs& attrs) { |
| return std::vector<ResourceRequest>{ResourceRequest::kTempSpace}; |
| }) |
| .set_attr<nnvm::TIsBackward>("TIsBackward", true) |
| .set_attr<FCompute>("FCompute<cpu>", LaOpBackward<cpu, 2, 2, 2, 1, sumlogdiag_backward>); |
| |
| NNVM_REGISTER_OP(_linalg_extractdiag) |
| .add_alias("linalg_extractdiag") |
| .describe(R"code(Extracts the diagonal entries of a square matrix. |
| Input is a tensor *A* of dimension *n >= 2*. |
| |
| If *n=2*, then *A* represents a single square matrix which diagonal elements get extracted as a 1-dimensional tensor. |
| |
| If *n>2*, then *A* represents a batch of square matrices on the trailing two dimensions. The extracted diagonals are returned as an *n-1*-dimensional tensor. |
| |
| .. note:: The operator supports float32 and float64 data types only. |
| |
| Examples:: |
| |
| Single matrix diagonal extraction |
| A = [[1.0, 2.0], |
| [3.0, 4.0]] |
| |
| extractdiag(A) = [1.0, 4.0] |
| |
| extractdiag(A, 1) = [2.0] |
| |
| Batch matrix diagonal extraction |
| A = [[[1.0, 2.0], |
| [3.0, 4.0]], |
| [[5.0, 6.0], |
| [7.0, 8.0]]] |
| |
| extractdiag(A) = [[1.0, 4.0], |
| [5.0, 8.0]] |
| )code" ADD_FILELINE) |
| .set_num_inputs(1) |
| .set_num_outputs(1) |
| .set_attr_parser(ParamParser<LaDiagParam>) |
| .set_attr<nnvm::FListInputNames>("FListInputNames", |
| [](const NodeAttrs& attrs) { |
| return std::vector<std::string>{"A"}; |
| }) |
| .set_attr<mxnet::FInferShape>("FInferShape", LaDiagTrianShape<true, true>) |
| .set_attr<nnvm::FInferType>("FInferType", ElemwiseType<1, 1>) |
| .set_attr<FCompute>("FCompute<cpu>", LaOpForward<cpu, 2, 1, 1, 1, copydiag>) |
| .set_attr<nnvm::FGradient>("FGradient", ElemwiseGradUseNone{"_backward_linalg_extractdiag"}) |
| .add_argument("A", "NDArray-or-Symbol", "Tensor of square matrices") |
| .add_arguments(LaDiagParam::__FIELDS__()); |
| |
| NNVM_REGISTER_OP(_backward_linalg_extractdiag) |
| .set_num_inputs(1) |
| .set_num_outputs(1) |
| .set_attr_parser(ParamParser<LaDiagParam>) |
| .set_attr<FResourceRequest>("FResourceRequest", |
| [](const NodeAttrs& attrs) { |
| return std::vector<ResourceRequest>{ResourceRequest::kTempSpace}; |
| }) |
| .set_attr<nnvm::TIsBackward>("TIsBackward", true) |
| .set_attr<FCompute>("FCompute<cpu>", LaOpBackward<cpu, 1, 2, 1, 1, copydiag>); |
| |
| NNVM_REGISTER_OP(_linalg_makediag) |
| .add_alias("linalg_makediag") |
| .describe(R"code(Constructs a square matrix with the input as diagonal. |
| Input is a tensor *A* of dimension *n >= 1*. |
| |
| If *n=1*, then *A* represents the diagonal entries of a single square matrix. This matrix will be returned as a 2-dimensional tensor. |
| If *n>1*, then *A* represents a batch of diagonals of square matrices. The batch of diagonal matrices will be returned as an *n+1*-dimensional tensor. |
| |
| .. note:: The operator supports float32 and float64 data types only. |
| |
| Examples:: |
| |
| Single diagonal matrix construction |
| A = [1.0, 2.0] |
| |
| makediag(A) = [[1.0, 0.0], |
| [0.0, 2.0]] |
| |
| makediag(A, 1) = [[0.0, 1.0, 0.0], |
| [0.0, 0.0, 2.0], |
| [0.0, 0.0, 0.0]] |
| |
| Batch diagonal matrix construction |
| A = [[1.0, 2.0], |
| [3.0, 4.0]] |
| |
| makediag(A) = [[[1.0, 0.0], |
| [0.0, 2.0]], |
| [[3.0, 0.0], |
| [0.0, 4.0]]] |
| )code" ADD_FILELINE) |
| .set_num_inputs(1) |
| .set_num_outputs(1) |
| .set_attr_parser(ParamParser<LaDiagParam>) |
| .set_attr<nnvm::FListInputNames>("FListInputNames", |
| [](const NodeAttrs& attrs) { |
| return std::vector<std::string>{"A"}; |
| }) |
| .set_attr<mxnet::FInferShape>("FInferShape", LaDiagTrianShape<true, false>) |
| .set_attr<nnvm::FInferType>("FInferType", ElemwiseType<1, 1>) |
| .set_attr<FCompute>("FCompute<cpu>", LaOpForward<cpu, 1, 2, 1, 1, copydiag>) |
| .set_attr<nnvm::FGradient>("FGradient", ElemwiseGradUseNone{"_backward_linalg_makediag"}) |
| .add_argument("A", "NDArray-or-Symbol", "Tensor of diagonal entries") |
| .add_arguments(LaDiagParam::__FIELDS__()); |
| |
| NNVM_REGISTER_OP(_backward_linalg_makediag) |
| .set_num_inputs(1) |
| .set_num_outputs(1) |
| .set_attr_parser(ParamParser<LaDiagParam>) |
| .set_attr<FResourceRequest>("FResourceRequest", |
| [](const NodeAttrs& attrs) { |
| return std::vector<ResourceRequest>{ResourceRequest::kTempSpace}; |
| }) |
| .set_attr<nnvm::TIsBackward>("TIsBackward", true) |
| .set_attr<FCompute>("FCompute<cpu>", LaOpBackward<cpu, 2, 1, 1, 1, copydiag>); |
| |
| NNVM_REGISTER_OP(_linalg_extracttrian) |
| .add_alias("linalg_extracttrian") |
| .describe(R"code(Extracts a triangular sub-matrix from a square matrix. |
| Input is a tensor *A* of dimension *n >= 2*. |
| |
| If *n=2*, then *A* represents a single square matrix from which a triangular sub-matrix is extracted as a 1-dimensional tensor. |
| |
| If *n>2*, then *A* represents a batch of square matrices on the trailing two dimensions. The extracted triangular sub-matrices are returned as an *n-1*-dimensional tensor. |
| |
| The *offset* and *lower* parameters determine the triangle to be extracted: |
| |
| - When *offset = 0* either the lower or upper triangle with respect to the main diagonal is extracted depending on the value of parameter *lower*. |
| - When *offset = k > 0* the upper triangle with respect to the k-th diagonal above the main diagonal is extracted. |
| - When *offset = k < 0* the lower triangle with respect to the k-th diagonal below the main diagonal is extracted. |
| |
| .. note:: The operator supports float32 and float64 data types only. |
| |
| Examples:: |
| |
| Single triagonal extraction |
| A = [[1.0, 2.0], |
| [3.0, 4.0]] |
| |
| extracttrian(A) = [1.0, 3.0, 4.0] |
| extracttrian(A, lower=False) = [1.0, 2.0, 4.0] |
| extracttrian(A, 1) = [2.0] |
| extracttrian(A, -1) = [3.0] |
| |
| Batch triagonal extraction |
| A = [[[1.0, 2.0], |
| [3.0, 4.0]], |
| [[5.0, 6.0], |
| [7.0, 8.0]]] |
| |
| extracttrian(A) = [[1.0, 3.0, 4.0], |
| [5.0, 7.0, 8.0]] |
| )code" ADD_FILELINE) |
| .set_num_inputs(1) |
| .set_num_outputs(1) |
| .set_attr_parser(ParamParser<LaTrianParam>) |
| .set_attr<nnvm::FListInputNames>("FListInputNames", |
| [](const NodeAttrs& attrs) { |
| return std::vector<std::string>{"A"}; |
| }) |
| .set_attr<mxnet::FInferShape>("FInferShape", LaDiagTrianShape<false, true>) |
| .set_attr<nnvm::FInferType>("FInferType", ElemwiseType<1, 1>) |
| .set_attr<FCompute>("FCompute<cpu>", LaOpForward<cpu, 2, 1, 1, 1, copytrian>) |
| .set_attr<nnvm::FGradient>("FGradient", ElemwiseGradUseNone{"_backward_linalg_extracttrian"}) |
| .add_argument("A", "NDArray-or-Symbol", "Tensor of square matrices") |
| .add_arguments(LaTrianParam::__FIELDS__()); |
| |
| NNVM_REGISTER_OP(_backward_linalg_extracttrian) |
| .set_num_inputs(1) |
| .set_num_outputs(1) |
| .set_attr_parser(ParamParser<LaTrianParam>) |
| .set_attr<FResourceRequest>("FResourceRequest", |
| [](const NodeAttrs& attrs) { |
| return std::vector<ResourceRequest>{ResourceRequest::kTempSpace}; |
| }) |
| .set_attr<nnvm::TIsBackward>("TIsBackward", true) |
| .set_attr<FCompute>("FCompute<cpu>", LaOpBackward<cpu, 1, 2, 1, 1, copytrian>); |
| |
| NNVM_REGISTER_OP(_linalg_maketrian) |
| .add_alias("linalg_maketrian") |
| .describe( |
| R"code(Constructs a square matrix with the input representing a specific triangular sub-matrix. |
| This is basically the inverse of *linalg.extracttrian*. Input is a tensor *A* of dimension *n >= 1*. |
| |
| If *n=1*, then *A* represents the entries of a triangular matrix which is lower triangular if *offset<0* or *offset=0*, *lower=true*. The resulting matrix is derived by first constructing the square |
| matrix with the entries outside the triangle set to zero and then adding *offset*-times an additional |
| diagonal with zero entries to the square matrix. |
| |
| If *n>1*, then *A* represents a batch of triangular sub-matrices. The batch of corresponding square matrices is returned as an *n+1*-dimensional tensor. |
| |
| .. note:: The operator supports float32 and float64 data types only. |
| |
| Examples:: |
| |
| Single matrix construction |
| A = [1.0, 2.0, 3.0] |
| |
| maketrian(A) = [[1.0, 0.0], |
| [2.0, 3.0]] |
| |
| maketrian(A, lower=false) = [[1.0, 2.0], |
| [0.0, 3.0]] |
| |
| maketrian(A, offset=1) = [[0.0, 1.0, 2.0], |
| [0.0, 0.0, 3.0], |
| [0.0, 0.0, 0.0]] |
| maketrian(A, offset=-1) = [[0.0, 0.0, 0.0], |
| [1.0, 0.0, 0.0], |
| [2.0, 3.0, 0.0]] |
| |
| Batch matrix construction |
| A = [[1.0, 2.0, 3.0], |
| [4.0, 5.0, 6.0]] |
| |
| maketrian(A) = [[[1.0, 0.0], |
| [2.0, 3.0]], |
| [[4.0, 0.0], |
| [5.0, 6.0]]] |
| |
| maketrian(A, offset=1) = [[[0.0, 1.0, 2.0], |
| [0.0, 0.0, 3.0], |
| [0.0, 0.0, 0.0]], |
| [[0.0, 4.0, 5.0], |
| [0.0, 0.0, 6.0], |
| [0.0, 0.0, 0.0]]] |
| )code" ADD_FILELINE) |
| .set_num_inputs(1) |
| .set_num_outputs(1) |
| .set_attr_parser(ParamParser<LaTrianParam>) |
| .set_attr<nnvm::FListInputNames>("FListInputNames", |
| [](const NodeAttrs& attrs) { |
| return std::vector<std::string>{"A"}; |
| }) |
| .set_attr<mxnet::FInferShape>("FInferShape", LaDiagTrianShape<false, false>) |
| .set_attr<nnvm::FInferType>("FInferType", ElemwiseType<1, 1>) |
| .set_attr<FCompute>("FCompute<cpu>", LaOpForward<cpu, 1, 2, 1, 1, copytrian>) |
| .set_attr<nnvm::FGradient>("FGradient", ElemwiseGradUseNone{"_backward_linalg_maketrian"}) |
| .add_argument("A", "NDArray-or-Symbol", "Tensor of triangular matrices stored as vectors") |
| .add_arguments(LaTrianParam::__FIELDS__()); |
| |
| NNVM_REGISTER_OP(_backward_linalg_maketrian) |
| .set_num_inputs(1) |
| .set_num_outputs(1) |
| .set_attr_parser(ParamParser<LaTrianParam>) |
| .set_attr<FResourceRequest>("FResourceRequest", |
| [](const NodeAttrs& attrs) { |
| return std::vector<ResourceRequest>{ResourceRequest::kTempSpace}; |
| }) |
| .set_attr<nnvm::TIsBackward>("TIsBackward", true) |
| .set_attr<FCompute>("FCompute<cpu>", LaOpBackward<cpu, 2, 1, 1, 1, copytrian>); |
| |
| NNVM_REGISTER_OP(_linalg_syrk) |
| .add_alias("linalg_syrk") |
| .describe(R"code(Multiplication of matrix with its transpose. |
| Input is a tensor *A* of dimension *n >= 2*. |
| |
| If *n=2*, the operator performs the BLAS3 function *syrk*: |
| |
| *out* = *alpha* \* *A* \* *A*\ :sup:`T` |
| |
| if *transpose=False*, or |
| |
| *out* = *alpha* \* *A*\ :sup:`T` \ \* *A* |
| |
| if *transpose=True*. |
| |
| If *n>2*, *syrk* is performed separately on the trailing two dimensions for all |
| inputs (batch mode). |
| |
| .. note:: The operator supports float32 and float64 data types only. |
| |
| Examples:: |
| |
| Single matrix multiply |
| A = [[1., 2., 3.], [4., 5., 6.]] |
| syrk(A, alpha=1., transpose=False) |
| = [[14., 32.], |
| [32., 77.]] |
| syrk(A, alpha=1., transpose=True) |
| = [[17., 22., 27.], |
| [22., 29., 36.], |
| [27., 36., 45.]] |
| |
| Batch matrix multiply |
| A = [[[1., 1.]], [[0.1, 0.1]]] |
| syrk(A, alpha=2., transpose=False) = [[[4.]], [[0.04]]] |
| )code" ADD_FILELINE) |
| .set_num_inputs(1) |
| .set_num_outputs(1) |
| .set_attr_parser(ParamParser<LaSyrkParam>) |
| .set_attr<nnvm::FListInputNames>("FListInputNames", |
| [](const NodeAttrs& attrs) { |
| return std::vector<std::string>{"A"}; |
| }) |
| .set_attr<mxnet::FInferShape>("FInferShape", LaSyrkShape) |
| .set_attr<nnvm::FInferType>("FInferType", ElemwiseType<1, 1>) |
| .set_attr<FCompute>("FCompute<cpu>", LaOpForward<cpu, 2, 2, 1, 1, syrk>) |
| .set_attr<nnvm::FGradient>("FGradient", ElemwiseGradUseIn{"_backward_linalg_syrk"}) |
| .add_argument("A", "NDArray-or-Symbol", "Tensor of input matrices") |
| .add_arguments(LaSyrkParam::__FIELDS__()); |
| |
| NNVM_REGISTER_OP(_backward_linalg_syrk) |
| .set_num_inputs(2) |
| .set_num_outputs(1) |
| .set_attr_parser(ParamParser<LaSyrkParam>) |
| .set_attr<FResourceRequest>("FResourceRequest", |
| [](const NodeAttrs& attrs) { |
| return std::vector<ResourceRequest>{ResourceRequest::kTempSpace}; |
| }) |
| .set_attr<nnvm::TIsBackward>("TIsBackward", true) |
| .set_attr<FCompute>("FCompute<cpu>", LaOpBackward<cpu, 2, 2, 2, 1, syrk_backward>); |
| |
| NNVM_REGISTER_OP(_linalg_gelqf) |
| .add_alias("linalg_gelqf") |
| .describe(R"code(LQ factorization for general matrix. |
| Input is a tensor *A* of dimension *n >= 2*. |
| |
| If *n=2*, we compute the LQ factorization (LAPACK *gelqf*, followed by *orglq*). *A* |
| must have shape *(x, y)* with *x <= y*, and must have full rank *=x*. The LQ |
| factorization consists of *L* with shape *(x, x)* and *Q* with shape *(x, y)*, so |
| that: |
| |
| *A* = *L* \* *Q* |
| |
| Here, *L* is lower triangular (upper triangle equal to zero) with nonzero diagonal, |
| and *Q* is row-orthonormal, meaning that |
| |
| *Q* \* *Q*\ :sup:`T` |
| |
| is equal to the identity matrix of shape *(x, x)*. |
| |
| If *n>2*, *gelqf* is performed separately on the trailing two dimensions for all |
| inputs (batch mode). |
| |
| .. note:: The operator supports float32 and float64 data types only. |
| |
| Examples:: |
| |
| Single LQ factorization |
| A = [[1., 2., 3.], [4., 5., 6.]] |
| Q, L = gelqf(A) |
| Q = [[-0.26726124, -0.53452248, -0.80178373], |
| [0.87287156, 0.21821789, -0.43643578]] |
| L = [[-3.74165739, 0.], |
| [-8.55235974, 1.96396101]] |
| |
| Batch LQ factorization |
| A = [[[1., 2., 3.], [4., 5., 6.]], |
| [[7., 8., 9.], [10., 11., 12.]]] |
| Q, L = gelqf(A) |
| Q = [[[-0.26726124, -0.53452248, -0.80178373], |
| [0.87287156, 0.21821789, -0.43643578]], |
| [[-0.50257071, -0.57436653, -0.64616234], |
| [0.7620735, 0.05862104, -0.64483142]]] |
| L = [[[-3.74165739, 0.], |
| [-8.55235974, 1.96396101]], |
| [[-13.92838828, 0.], |
| [-19.09768702, 0.52758934]]] |
| )code" ADD_FILELINE) |
| .set_num_inputs(1) |
| .set_num_outputs(2) |
| .set_attr<nnvm::FListInputNames>("FListInputNames", |
| [](const NodeAttrs& attrs) { |
| return std::vector<std::string>{"A"}; |
| }) |
| .set_attr<mxnet::FInferShape>("FInferShape", LaLQFactShape) |
| .set_attr<nnvm::FInferType>("FInferType", ElemwiseType<1, 2>) |
| .set_attr<nnvm::FInplaceOption>("FInplaceOption", |
| [](const NodeAttrs& attrs) { |
| return std::vector<std::pair<int, int>>{{0, 0}}; |
| }) |
| .set_attr<FResourceRequest>("FResourceRequest", |
| [](const NodeAttrs& attrs) { |
| return std::vector<ResourceRequest>{ResourceRequest::kTempSpace}; |
| }) |
| .set_attr<THasDeterministicOutput>("THasDeterministicOutput", true) |
| .set_attr<FCompute>("FCompute<cpu>", LaOpForward<cpu, 2, 2, 1, 2, gelqf>) |
| .set_attr<nnvm::FGradient>("FGradient", ElemwiseGradUseOut{"_backward_linalg_gelqf"}) |
| .add_argument("A", "NDArray-or-Symbol", "Tensor of input matrices to be factorized"); |
| |
| NNVM_REGISTER_OP(_backward_linalg_gelqf) |
| .set_num_inputs(4) |
| .set_num_outputs(1) |
| .set_attr<nnvm::FInplaceOption>("FInplaceOption", |
| [](const NodeAttrs& attrs) { |
| return std::vector<std::pair<int, int>>{{0, 0}}; |
| }) |
| .set_attr<FResourceRequest>("FResourceRequest", |
| [](const NodeAttrs& attrs) { |
| return std::vector<ResourceRequest>{ResourceRequest::kTempSpace}; |
| }) |
| .set_attr<nnvm::TIsBackward>("TIsBackward", true) |
| .set_attr<FCompute>("FCompute<cpu>", LaOpBackward<cpu, 2, 2, 4, 1, gelqf_backward>); |
| |
| NNVM_REGISTER_OP(_linalg_syevd) |
| .describe(R"code(Eigendecomposition for symmetric matrix. |
| Input is a tensor *A* of dimension *n >= 2*. |
| |
| If *n=2*, *A* must be symmetric, of shape *(x, x)*. We compute the eigendecomposition, |
| resulting in the orthonormal matrix *U* of eigenvectors, shape *(x, x)*, and the |
| vector *L* of eigenvalues, shape *(x,)*, so that: |
| |
| *U* \* *A* = *diag(L)* \* *U* |
| |
| Here: |
| |
| *U* \* *U*\ :sup:`T` = *U*\ :sup:`T` \* *U* = *I* |
| |
| where *I* is the identity matrix. Also, *L(0) <= L(1) <= L(2) <= ...* (ascending order). |
| |
| If *n>2*, *syevd* is performed separately on the trailing two dimensions of *A* (batch |
| mode). In this case, *U* has *n* dimensions like *A*, and *L* has *n-1* dimensions. |
| |
| .. note:: The operator supports float32 and float64 data types only. |
| |
| .. note:: Derivatives for this operator are defined only if *A* is such that all its |
| eigenvalues are distinct, and the eigengaps are not too small. If you need |
| gradients, do not apply this operator to matrices with multiple eigenvalues. |
| |
| Examples:: |
| |
| Single symmetric eigendecomposition |
| A = [[1., 2.], [2., 4.]] |
| U, L = syevd(A) |
| U = [[0.89442719, -0.4472136], |
| [0.4472136, 0.89442719]] |
| L = [0., 5.] |
| |
| Batch symmetric eigendecomposition |
| A = [[[1., 2.], [2., 4.]], |
| [[1., 2.], [2., 5.]]] |
| U, L = syevd(A) |
| U = [[[0.89442719, -0.4472136], |
| [0.4472136, 0.89442719]], |
| [[0.92387953, -0.38268343], |
| [0.38268343, 0.92387953]]] |
| L = [[0., 5.], |
| [0.17157288, 5.82842712]] |
| )code" ADD_FILELINE) |
| .set_num_inputs(1) |
| .set_num_outputs(2) |
| .set_attr<nnvm::FListInputNames>("FListInputNames", |
| [](const NodeAttrs& attrs) { |
| return std::vector<std::string>{"A"}; |
| }) |
| .set_attr<mxnet::FInferShape>("FInferShape", LaEigFactShape) |
| .set_attr<nnvm::FInferType>("FInferType", ElemwiseType<1, 2>) |
| .set_attr<nnvm::FInplaceOption>("FInplaceOption", |
| [](const NodeAttrs& attrs) { |
| return std::vector<std::pair<int, int>>{{0, 0}}; |
| }) |
| .set_attr<FResourceRequest>("FResourceRequest", |
| [](const NodeAttrs& attrs) { |
| return std::vector<ResourceRequest>{ResourceRequest::kTempSpace}; |
| }) |
| .set_attr<THasDeterministicOutput>("THasDeterministicOutput", true) |
| .set_attr<FCompute>("FCompute<cpu>", LaOpForwSyevd<cpu, syevd>) |
| .set_attr<nnvm::FGradient>("FGradient", ElemwiseGradUseOut{"_backward_linalg_syevd"}) |
| .add_argument("A", "NDArray-or-Symbol", "Tensor of input matrices to be factorized"); |
| |
| NNVM_REGISTER_OP(_backward_linalg_syevd) |
| .set_num_inputs(4) |
| .set_num_outputs(1) |
| .set_attr<nnvm::FInplaceOption>("FInplaceOption", |
| [](const NodeAttrs& attrs) { |
| return std::vector<std::pair<int, int>>{{0, 0}}; |
| }) |
| .set_attr<FResourceRequest>("FResourceRequest", |
| [](const NodeAttrs& attrs) { |
| return std::vector<ResourceRequest>{ResourceRequest::kTempSpace}; |
| }) |
| .set_attr<nnvm::TIsBackward>("TIsBackward", true) |
| .set_attr<FCompute>("FCompute<cpu>", LaOpBackwSyevd<cpu, syevd_backward>); |
| |
| NNVM_REGISTER_OP(_linalg_inverse) |
| .add_alias("linalg_inverse") |
| .add_alias("_npi_inv") |
| .describe(R"code(Compute the inverse of a matrix. |
| Input is a tensor *A* of dimension *n >= 2*. |
| |
| If *n=2*, *A* is a square matrix. We compute: |
| |
| *out* = *A*\ :sup:`-1` |
| |
| If *n>2*, *inverse* is performed separately on the trailing two dimensions |
| for all inputs (batch mode). |
| |
| .. note:: The operator supports float32 and float64 data types only. |
| |
| Examples:: |
| |
| Single matrix inverse |
| A = [[1., 4.], [2., 3.]] |
| inverse(A) = [[-0.6, 0.8], [0.4, -0.2]] |
| |
| Batch matrix inverse |
| A = [[[1., 4.], [2., 3.]], |
| [[1., 3.], [2., 4.]]] |
| inverse(A) = [[[-0.6, 0.8], [0.4, -0.2]], |
| [[-2., 1.5], [1., -0.5]]] |
| )code" ADD_FILELINE) |
| .set_num_inputs(1) |
| .set_num_outputs(1) |
| .set_attr<nnvm::FListInputNames>("FListInputNames", |
| [](const NodeAttrs& attrs) { |
| return std::vector<std::string>{"A"}; |
| }) |
| .set_attr<mxnet::FInferShape>("FInferShape", InverseShape) |
| .set_attr<nnvm::FInferType>("FInferType", ElemwiseType<1, 1>) |
| .set_attr<nnvm::FInplaceOption>("FInplaceOption", |
| [](const NodeAttrs& attrs) { |
| return std::vector<std::pair<int, int>>{{0, 0}}; |
| }) |
| .set_attr<FResourceRequest>("FResourceRequest", |
| [](const NodeAttrs& attrs) { |
| return std::vector<ResourceRequest>{ResourceRequest::kTempSpace}; |
| }) |
| .set_attr<THasDeterministicOutput>("THasDeterministicOutput", true) |
| .set_attr<FCompute>("FCompute<cpu>", LaOpForward<cpu, 2, 2, 1, 1, inverse>) |
| .set_attr<nnvm::FGradient>("FGradient", ElemwiseGradUseOut{"_backward_linalg_inverse"}) |
| .add_argument("A", "NDArray-or-Symbol", "Tensor of square matrix"); |
| |
| NNVM_REGISTER_OP(_backward_linalg_inverse) |
| .set_num_inputs(2) |
| .set_num_outputs(1) |
| .set_attr<nnvm::FInplaceOption>("FInplaceOption", |
| [](const NodeAttrs& attrs) { |
| return std::vector<std::pair<int, int>>{{0, 0}}; |
| }) |
| .set_attr<FResourceRequest>("FResourceRequest", |
| [](const NodeAttrs& attrs) { |
| return std::vector<ResourceRequest>{ResourceRequest::kTempSpace}; |
| }) |
| .set_attr<nnvm::TIsBackward>("TIsBackward", true) |
| .set_attr<FCompute>("FCompute<cpu>", LaOpBackward<cpu, 2, 2, 2, 1, inverse_backward>); |
| |
| NNVM_REGISTER_OP(_linalg_det) |
| .add_alias("linalg_det") |
| .add_alias("_npi_det") |
| .describe(R"code(Compute the determinant of a matrix. |
| Input is a tensor *A* of dimension *n >= 2*. |
| |
| If *n=2*, *A* is a square matrix. We compute: |
| |
| *out* = *det(A)* |
| |
| If *n>2*, *det* is performed separately on the trailing two dimensions |
| for all inputs (batch mode). |
| |
| .. note:: The operator supports float32 and float64 data types only. |
| .. note:: There is no gradient backwarded when A is non-invertible (which is |
| equivalent to det(A) = 0) because zero is rarely hit upon in float |
| point computation and the Jacobi's formula on determinant gradient |
| is not computationally efficient when A is non-invertible. |
| |
| Examples:: |
| |
| Single matrix determinant |
| A = [[1., 4.], [2., 3.]] |
| det(A) = [-5.] |
| |
| Batch matrix determinant |
| A = [[[1., 4.], [2., 3.]], |
| [[2., 3.], [1., 4.]]] |
| det(A) = [-5., 5.] |
| )code" ADD_FILELINE) |
| .set_num_inputs(1) |
| .set_num_outputs(3) |
| .set_attr<nnvm::FListInputNames>("FListInputNames", |
| [](const NodeAttrs& attrs) { |
| return std::vector<std::string>{"A"}; |
| }) |
| .set_attr<nnvm::FNumVisibleOutputs>("FNumVisibleOutputs", |
| [](const NodeAttrs& attrs) { return 1; }) |
| .set_attr<mxnet::FInferShape>("FInferShape", DetShape<1>) |
| .set_attr<nnvm::FInferType>("FInferType", DetType<1>) |
| .set_attr<FResourceRequest>("FResourceRequest", |
| [](const NodeAttrs& attrs) { |
| return std::vector<ResourceRequest>{ResourceRequest::kTempSpace}; |
| }) |
| .set_attr<THasDeterministicOutput>("THasDeterministicOutput", true) |
| .set_attr<FCompute>("FCompute<cpu>", LaOpDetForward<cpu, 1, det>) |
| .set_attr<nnvm::FGradient>("FGradient", ReduceDetGrad<1>{"_backward_linalg_det"}) |
| .add_argument("A", "NDArray-or-Symbol", "Tensor of square matrix"); |
| |
| NNVM_REGISTER_OP(_backward_linalg_det) |
| .set_num_inputs(4) |
| .set_num_outputs(1) |
| .set_attr<FResourceRequest>("FResourceRequest", |
| [](const NodeAttrs& attrs) { |
| return std::vector<ResourceRequest>{ResourceRequest::kTempSpace}; |
| }) |
| .set_attr<nnvm::TIsBackward>("TIsBackward", true) |
| .set_attr<FCompute>("FCompute<cpu>", LaOpDetBackward<cpu, 1, det_backward>); |
| |
| NNVM_REGISTER_OP(_linalg_slogdet) |
| .add_alias("linalg_slogdet") |
| .add_alias("_npi_slogdet") |
| .describe(R"code(Compute the sign and log of the determinant of a matrix. |
| Input is a tensor *A* of dimension *n >= 2*. |
| |
| If *n=2*, *A* is a square matrix. We compute: |
| |
| *sign* = *sign(det(A))* |
| *logabsdet* = *log(abs(det(A)))* |
| |
| If *n>2*, *slogdet* is performed separately on the trailing two dimensions |
| for all inputs (batch mode). |
| |
| .. note:: The operator supports float32 and float64 data types only. |
| .. note:: The gradient is not properly defined on sign, so the gradient of |
| it is not backwarded. |
| .. note:: No gradient is backwarded when A is non-invertible. Please see |
| the docs of operator det for detail. |
| |
| Examples:: |
| |
| Single matrix signed log determinant |
| A = [[2., 3.], [1., 4.]] |
| sign, logabsdet = slogdet(A) |
| sign = [1.] |
| logabsdet = [1.609438] |
| |
| Batch matrix signed log determinant |
| A = [[[2., 3.], [1., 4.]], |
| [[1., 2.], [2., 4.]], |
| [[1., 2.], [4., 3.]]] |
| sign, logabsdet = slogdet(A) |
| sign = [1., 0., -1.] |
| logabsdet = [1.609438, -inf, 1.609438] |
| )code" ADD_FILELINE) |
| .set_num_inputs(1) |
| .set_num_outputs(4) |
| .set_attr<nnvm::FListInputNames>("FListInputNames", |
| [](const NodeAttrs& attrs) { |
| return std::vector<std::string>{"A"}; |
| }) |
| .set_attr<nnvm::FNumVisibleOutputs>("FNumVisibleOutputs", |
| [](const NodeAttrs& attrs) { return 2; }) |
| .set_attr<mxnet::FInferShape>("FInferShape", DetShape<2>) |
| .set_attr<nnvm::FInferType>("FInferType", DetType<2>) |
| .set_attr<FResourceRequest>("FResourceRequest", |
| [](const NodeAttrs& attrs) { |
| return std::vector<ResourceRequest>{ResourceRequest::kTempSpace}; |
| }) |
| .set_attr<FCompute>("FCompute<cpu>", LaOpDetForward<cpu, 2, slogdet>) |
| .set_attr<nnvm::FGradient>("FGradient", ReduceDetGrad<2>{"_backward_linalg_slogdet"}) |
| .add_argument("A", "NDArray-or-Symbol", "Tensor of square matrix"); |
| |
| NNVM_REGISTER_OP(_backward_linalg_slogdet) |
| .set_num_inputs(5) |
| .set_num_outputs(1) |
| .set_attr<FResourceRequest>("FResourceRequest", |
| [](const NodeAttrs& attrs) { |
| return std::vector<ResourceRequest>{ResourceRequest::kTempSpace}; |
| }) |
| .set_attr<nnvm::TIsBackward>("TIsBackward", true) |
| .set_attr<FCompute>("FCompute<cpu>", LaOpDetBackward<cpu, 2, slogdet_backward>); |
| |
| } // namespace op |
| } // namespace mxnet |
