MXNet ndarray.linalg 线性代数算子 API 全解析：从 GEMM 到特征分解的完整指南

MXNet ndarray.linalg 线性代数算子 API 全解析从 GEMM 到特征分解的完整指南【免费下载链接】mxnetLightweight, Portable, Flexible Distributed/Mobile Deep Learning with Dynamic, Mutation-aware Dataflow Dep Scheduler; for Python, R, Julia, Scala, Go, Javascript and more项目地址: https://gitcode.com/gh_mirrors/mx/mxnet本文是 MXNet 旧式legacy线性代数算子模块mxnet.ndarray.linalg的深度技术指南。它以文档页 docs/python_docs/python/api/legacy/ndarray/linalg/index.rstSphinxautomodule/autosummary自动生成的 API 参考为核心骨架结合src/operator/tensor/la_op.cc中算子注册与完整 docstring系统讲解 17 个线性代数算子gemm、potrf、trsm、syevd、det 等的数学定义、参数默认值、批处理语义与梯度支持。读完本文你将掌握mx.nd.linalg与mx.sym.linalg两套 API 的完整用法并理解其底层如何调用 BLAS/LAPACK 及如何参与自动求导。一、模块定位文档页背后是什么index.rst正文只有一行指令.. automodule:: mxnet.ndarray.linalg :members: :autosummary:它本身不承载文字内容而是告诉 Sphinx 从mxnet.ndarray.linalg模块的 docstring 中自动生成完整 API 参考。因此这份文档的实体内容实际分布在两处Python 模块入口python/mxnet/ndarray/linalg.py 通过from .gen_linalg import *导入构建期自动生成的算子绑定gen_linalg为生成产物构建后才会出现同名符号 API 见 python/mxnet/symbol/linalg.py。模块在 python/mxnet/ndarray/init.py 中以from . import ... linalg ...注册并列入__all__。C 算子注册与 docstringsrc/operator/tensor/la_op.cc 中每个算子通过NNVM_REGISTER_OP注册.describe(...)里的Rcode(...)字符串就是渲染进 API 文档的核心内容。值得区分的是mx.nd.linalg本页文档主题是旧式 NDArray API与之并行的还有 NumPy 兼容的新式np.linalg见 python/mxnet/ndarray/numpy/linalg.py两者算子底层有部分复用如_linalg_inverse同时带_npi_inv别名_linalg_det带_npi_det别名。二、全局约定所有算子共享的语义从 la_op.cc 的 docstring 与 la_op-inl.h 的实现可以看出本模块 17 个算子遵循完全一致的约定约定说明输入维度输入张量维度n 2n 2表示单个矩阵n 2表示批处理模式批处理方式n 2时对所有输入在尾两维上逐矩阵独立执行运算la_op-inl.h中统一走linalg_batch_*系列见 la_op-inl.h数据类型仅支持 float32 与 float64每个算子 docstring 均有.. note::声明梯度支持每个前向算子都注册了对应的_backward_linalg_*反向算子FGradient属性可直接用于autograd原地计算部分算子声明了FInplaceOption例如 gemm 的{2, 0}表示输出可原地覆盖 C确定性输出gelqf、syevd、inverse、det、slogdet 均设置THasDeterministicOutput true临时空间多数反向算子请求ResourceRequest::kTempSpace作为工作区三、矩阵乘法类gemm 与 gemm23.1 linalg.gemm乘加一体注册代码见 la_op.cc。输入为三个张量 A、B、C维度n 2且前n-2维形状相同n 2时执行 BLAS3gemmout alpha * op(A) * op(B) beta * Cn 2时对尾两维逐矩阵批量执行参数表来自 LaMatrixMacParam见 la_op.h参数默认值含义transpose_afalse是否对 A 转置transpose_bfalse是否对 B 转置alpha1.0乘以 A·B 的标量因子beta1.0乘以 C 的标量因子axis-2矩阵行索引对应的轴见 3.3文档中的示例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_bTrue, alpha2.0, beta10.0) [[14.0, 14.0, 14.0], [14.0, 14.0, 14.0]]Python 调用import mxnet as mx A mx.nd.array([[1.0, 1.0], [1.0, 1.0]]) B mx.nd.array([[1.0, 1.0], [1.0, 1.0], [1.0, 1.0]]) C mx.nd.array([[1.0, 1.0, 1.0], [1.0, 1.0, 1.0]]) mx.nd.linalg.gemm(A, B, C, transpose_bTrue, alpha2.0, beta10.0)3.2 linalg.gemm2纯矩阵乘法注册代码见 la_op.cc。两个输入 A、Bout alpha * op(A) * op(B)无 beta/C。参数为transpose_a、transpose_b、alpha、axis默认同上。示例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_bTrue, alpha2.0) [[4.0, 4.0, 4.0], [4.0, 4.0, 4.0]]3.3 axis 参数矩阵编码轴默认axis -2即用尾两维编码矩阵倒数第二维为行、最后一维为列。当指定非默认axis时语义等价于swapaxes/gemm/swapaxes组合但无需额外的转置开销。docstring 给出的 5 维张量示例A1 swapaxes(A, dim11, dim23) B1 swapaxes(B, dim11, dim23) C swapaxes(C, dim11, dim23) C gemm(A1, B1, C) C swapaxis(C, dim11, dim23)形状推断在 la_op.h 的LaMatrixMultMacOpShape中完成要求两个输入除行/列轴外形状一致并校验乘法维度是否兼容非法axis会报错Invalid row axis。3.4 Tensor Core 加速提示当输入为 float32且环境变量MXNET_CUDA_ALLOW_TENSOR_CORE与MXNET_CUDA_TENSOR_OP_MATH_ALLOW_CONVERSION同时置为1时gemm/gemm2 会尝试以伪 float16 精度float32 计算、float16 读写使用 NVIDIA GPU 上的 Tensor Core可能带来显著加速。该行为在测试 tests/python/unittest/test_operator.py 的test_gemm中通过environment(MXNET_CUDA_TENSOR_OP_MATH_ALLOW_CONVERSION, ...)显式验证。四、Cholesky 分解族potrf 与 potri4.1 linalg.potrfCholesky 分解注册见 la_op.cc。输入 A 为对称正定矩阵n 2输出三角因子 B对角线为正lower trueA B * B^Tlower falseA B^T * B参数lower默认true见 LaCholeskyParamla_op.h。示例A [[4.0, 1.0], [1.0, 4.25]] potrf(A) [[2.0, 0], [0.5, 2.0]]4.2 linalg.potri由 Cholesky 因子求逆注册见 la_op.cc。输入为 Cholesky 因子 A三角、正对角lower trueout A^-T * A^-1lower falseout A^-1 * A^-T若 A 是 B 的 Cholesky 因子则out B^-1。docstring 特别警告只有在确实需要 B 的逆矩阵时才用本算子否则应使用 Cholesky 因子 Apotrf配合回代trsm后者数值上更安全且更便宜。五、三角矩阵运算trmm 与 trsm两个算子共用参数结构LaTriangMatrixMultParamla_op.h参数默认值含义transposefalse是否使用三角矩阵的转置rightsidefalse是否将三角矩阵乘在右侧lowertrue三角矩阵是下三角true还是上三角falsealpha1.0结果标量因子5.1 linalg.trmm三角矩阵乘法注册见 la_op.cc。A 为三角矩阵执行 BLAS3trmmrightsidefalseout alpha * op(A) * Brightsidetrueout alpha * B * op(A)示例A [[1.0, 0], [1.0, 1.0]] B [[1.0, 1.0, 1.0], [1.0, 1.0, 1.0]] trmm(A, B, alpha2.0) [[2.0, 2.0, 2.0], [4.0, 4.0, 4.0]]5.2 linalg.trsm三角回代求解注册见 la_op.cc。求解矩阵方程rightsidefalseop(A) * out alpha * Brightsidetrueout * op(A) alpha * B这是 4.2 中推荐的、数值上更安全的求解而非显式求逆路径。示例A [[1.0, 0], [1.0, 1.0]] B [[2.0, 2.0, 2.0], [4.0, 4.0, 4.0]] trsm(A, B, alpha0.5) [[1.0, 1.0, 1.0], [1.0, 1.0, 1.0]]六、对角线与三角抽取/构造extractdiag、makediag、extracttrian、maketrian6.1 linalg.extractdiag / linalg.makediag注册见 la_op.cc。二者共用参数offset默认0LaDiagParamla_op.h0为主对角线正值对应主对角线之上的第 k 条对角线负值对应之下。extractdiag(A)抽取方阵对角线n2时输出一维张量n2时输出n-1维张量makediag(A)由对角线构造方阵逆操作n1时输出二维矩阵n1时输出n1维张量A [[1.0, 2.0], [3.0, 4.0]] extractdiag(A) [1.0, 4.0] extractdiag(A, 1) [2.0] makediag([1.0, 2.0]) [[1.0, 0.0], [0.0, 2.0]] makediag([1.0, 2.0], 1) [[0.0, 1.0, 0.0], [0.0, 0.0, 2.0], [0.0, 0.0, 0.0]]6.2 linalg.extracttrian / linalg.maketrian注册见 la_op.cc。参数offset默认0与lower默认trueLaTrianParamla_op.hoffset 0按lower抽取/构造主对角线对应三角offset k 0主对角线之上第 k 条对角线对应的上三角offset k 0主对角线之下第 k 条对角线对应的下三角A [[1.0, 2.0], [3.0, 4.0]] extracttrian(A) [1.0, 3.0, 4.0] # 下三角 extracttrian(A, lowerFalse) [1.0, 2.0, 4.0] # 上三角 extracttrian(A, 1) [2.0] extracttrian(A, -1) [3.0] maketrian([1.0, 2.0, 3.0]) [[1.0, 0.0], [2.0, 3.0]] maketrian([1.0, 2.0, 3.0], offset1) [[0.0, 1.0, 2.0], [0.0, 0.0, 3.0], [0.0, 0.0, 0.0]]七、矩阵分解gelqfLQ与 syevd对称特征分解7.1 linalg.gelqfLQ 分解注册见 la_op.cc。底层调用 LAPACKgelqforglq。A 形状须为(x, y)且x y、满秩x输出两个张量Q形状(x, y)行正交Q * Q^T I与L形状(x, x)下三角、对角非零满足A L * Q。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]]7.2 linalg.syevd对称矩阵特征分解注册见 la_op.cc。A 为对称矩阵(x, x)输出正交特征向量矩阵UU * U^T I与特征值向量L升序排列L(0) L(1) ...满足U * A diag(L) * U。批处理时U保持与 A 相同维度L少一维。梯度注意事项docstring 明确警告仅当 A 的特征值全部互异且特征值间隔不太小时本算子的导数才被定义。需要对多重点矩阵求梯度时不要使用本算子。A [[1., 2.], [2., 4.]] U, L syevd(A) U [[0.89442719, -0.4472136], [0.4472136, 0.89442719]] L [0., 5.]八、求逆与行列式inverse、det、slogdet8.1 linalg.inverse注册见 la_op.cc。计算out A^-1方阵、批处理模式同上。示例A [[1., 4.], [2., 3.]] inverse(A) [[-0.6, 0.8], [0.4, -0.2]]8.2 linalg.det注册见 la_op.cc。计算行列式输出形状(1,)单矩阵或对应批维度。反向梯度通过ReduceDetGrad与_backward_linalg_det注册。梯度注意当 A 不可逆det(A) 0时不回传梯度——浮点计算中几乎不会精确命中 0且对不可逆矩阵使用 Jacobi 公式求行列式梯度并不高效。8.3 linalg.slogdet注册见 la_op.cc。输出两个张量sign sign(det(A))与logabsdet log(abs(det(A)))。两个梯度注意事项sign 的梯度定义不成立故不回传A 不可逆时同样不回传梯度。A [[2., 3.], [1., 4.]] sign, logabsdet slogdet(A) sign [1.] logabsdet [1.609438]九、辅助算子sumlogdiag 与 syrk9.1 linalg.sumlogdiag注册见 la_op.cc。计算方阵对角元素自然对数之和A 须为方阵且对角元素为正单矩阵输出形状(1,)批处理时对尾两维逐矩阵归约。典型用途是配合 Cholesky 分解计算log det(A) 2 * sumlogdiag(potrf(A))。A [[1.0, 1.0], [1.0, 7.0]] sumlogdiag(A) [1.9459]9.2 linalg.syrk注册见 la_op.cc。执行 BLAS3syrk矩阵与其转置之积transposefalseout alpha * A * A^Ttransposetrueout alpha * A^T * A参数alpha默认1.0LaSyrkParamla_op.h。A [[1., 2., 3.], [4., 5., 6.]] syrk(A, alpha1., transposeFalse) [[14., 32.], [32., 77.]] syrk(A, alpha1., transposeTrue) [[17., 22., 27.], [22., 29., 36.], [27., 36., 45.]]十、源码级原理算子如何注册与执行10.1 注册与自动生成每个算子以NNVM_REGISTER_OP(_linalg_xxx)注册并add_alias(linalg_xxx)因此mx.nd.linalg.gemm与内部名_linalg_gemm等价。算子同时声明FListInputNames如 gemm 的{A, B, C}、FInferShape、FInferType、FComputecpu、FGradient。Python 侧gen_linalg绑定由构建期脚本生成用户代码只需import mxnet后通过mx.nd.linalg.*/mx.sym.linalg.*访问。10.2 计算内核批处理模板与 BLAS/LAPACKla_op-inl.h 中的注释明确Forward computations (always using batched processing)所有算子统一走linalg_batch_gemm、linalg_batch_potrf、linalg_batch_trsm等批处理内核底层经由 src/operator/linalg.h 与 src/operator/linalg_impl.h 封装 BLAS/LAPACK 例程LAPACK 的 C 接口桥接层见 src/operator/c_lapack_api.cc 与 src/operator/c_lapack_api.h。GPU 侧对应实现位于 src/operator/tensor/la_op.cu。la_op-inl.h顶部还注明梯度计算细节可参考 arXiv:1710.08717。10.3 反向算子与自动求导每个前向算子都挂接FGradient指向_backward_linalg_*反向算子例如 gemm 用ElemwiseGradUseInpotrf 用ElemwiseGradUseOutinverse 用ElemwiseGradUseOut。因此这些算子可直接嵌入autograd计算图from mxnet import autograd, nd A nd.array([[4.0, 1.0], [1.0, 4.25]]) A.attach_grad() with autograd.record(): B nd.linalg.potrf(A) loss nd.sum(B) loss.backward() print(A.grad)十一、测试与验证单元测试集中在 tests/python/unittest/test_operator.py_gemm_test_helper约 L6019-L6115覆盖 gemm/gemm2 的四种转置组合transpose_a/transpose_b全排列、批处理通过rep_3x复制出(3, 1, m, n)批量、以及axis0、axis-3等非默认矩阵编码轴每个用例同时做前向数值校验check_symbolic_forward与数值梯度校验check_numeric_gradient并指定 rtol/atol 容差。test_gemmL6118-L6124在 float64 全量测试并在MXNET_CUDA_TENSOR_OP_MATH_ALLOW_CONVERSION置 0/1 两种环境下测试 float32验证 Tensor Core 路径开关。mx.nd.linalg.potrf还被 Gluon 概率分布模块如多元正态见 python/mxnet/gluon/probability/distributions/multivariate_normal.py在内部用于协方差分解测试中亦有对应验证见 test_operator.py。十二、快速参考算子一览Python APImx.nd.linalg/mx.sym.linalg数学语义关键参数默认值gemm(A, B, C)alpha·op(A)·op(B) beta·Ctranspose_afalse, transpose_bfalse, alpha1.0, beta1.0, axis-2gemm2(A, B)alpha·op(A)·op(B)transpose_afalse, transpose_bfalse, alpha1.0, axis-2potrf(A)CholeskyA B·B^T或B^T·Blowertruepotri(A)由 Cholesky 因子求逆lowertruetrmm(A, B)三角矩阵乘法transposefalse, rightsidefalse, lowertrue, alpha1.0trsm(A, B)三角回代求解同上sumlogdiag(A)对角元素自然对数之和无extractdiag(A)抽取对角线offset0makediag(A)由对角线构造方阵offset0extracttrian(A)抽取三角子矩阵offset0, lowertruemaketrian(A)由三角子矩阵构造方阵offset0, lowertruesyrk(A)alpha·A·A^T或alpha·A^T·Atransposefalse, alpha1.0gelqf(A)LQ 分解返回(Q, L)无syevd(A)对称特征分解返回(U, L)无inverse(A)矩阵求逆无det(A)行列式无slogdet(A)符号与行列式对数返回(sign, logabsdet)无使用要点回顾所有算子仅支持 float32/float64输入维度n 2且n 2时自动按尾两维批处理涉及对称/三角结构的算子务必按 docstring 约束传入合法矩阵如 potrf 要求对称正定、gelqf 要求x y且满秩、syevd 求梯度时要求特征值互异能用 trsm 回代求解时优先于显式求逆以保证数值稳定性与计算效率。【免费下载链接】mxnetLightweight, Portable, Flexible Distributed/Mobile Deep Learning with Dynamic, Mutation-aware Dataflow Dep Scheduler; for Python, R, Julia, Scala, Go, Javascript and more项目地址: https://gitcode.com/gh_mirrors/mx/mxnet创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考