jittor.linalg

这里是Jittor的线性代数函数的API文档,您可以通过from jittor import linalg来获取该模块。

基本函数简介

基本线性代数运算API

  • linalg.inv(a)

    对a进行求逆运算

  • linalg.pinv(a)

    对a进行广义求逆运算。该运算不要求原矩阵a可逆。

  • linalg.slogdet(a)

    对a求取slogdet。会返回值以及符号。

  • linalg.det(a)

    对a求行列式。

  • linalg.solve(a,b)

    求解线性方程Ax=b的解。

分解API

  • linalg.cholesky(a)

    对a进行cholesky分解。

  • linalg.qr(a)

    对a进行qr分解。

  • linalg.svd

    对a进行奇异值分解。

特征值API

  • linalg.eig(a)

    求取a的特征值以及特征向量。

  • linalg.eigh(a)

    针对埃尔米特矩阵或者对称矩阵求特征值以及特征向量。

目前的linalg库支持

jittor.linalg.cholesky(x)[源代码]

对输入矩阵的Cholesky分解, 形式如下公式:

\[x = LL^T\]

其中 \(x\) 必须是 Hermite 和正定矩阵。 \(L\) 是一个下三角矩阵。

参数:
  • x (Var): 输入的二维矩阵, 维度为 (..., M, M)\(x\) 应满足如下条件:首先, \(x\) 应是正定的;其次, \(x\) 应是 Hermite 的(即 \(x\) 等于其共轭转置)。 \(x\) 可以是实矩阵或者复矩阵。

返回值:

返回 \(x\) 的Cholesky分解的下三角矩阵 \(L\) ( Var ), 其维度与 \(x\) 相同, 为 (..., M, M)

代码示例:
>>> import jittor as jt
>>> a = jt.array([[4.0, 12, -16], [12, 37, -43], [-16, -43, 98]])
>>> print(jt.linalg.cholesky(a))
jt.Var([[ 2.  0.  0.]
        [ 6.  1.  0.]
        [-8.  5.  3.]], dtype=float32)
jittor.linalg.det(x)[源代码]

计算输入矩阵的行列式值。

参数:
  • x (Var): 输入的张量, 要求为形状 (..., M, M) 的张量, 其中 \(M\) 是矩阵的维度。此张量可以是任意维度, 最后两个维度被视为2D矩阵。

返回值:

输入矩阵 \(x\) 的行列式值, 形状为 (..., 1)Var

代码示例:
>>> import jittor as jt
>>> x = jt.array([[1.0,2.0],[3.0,4.0]])
>>> print(jt.linalg.det(x))
jt.Var([-2.], dtype=float32)
jittor.linalg.eigh(x)[源代码]

计算输入矩阵的特征值和特征向量。特征值和特征向量的计算公式:

\[A v = \lambda v\]

其中, \(A\) 是单位矩阵, \(v\) 是特征向量, \(\lambda\) 是特征值。

参数:
  • x (Var): 输入矩阵, 维度为(…, \(M\) , \(M\) ), 其中 \(M\) 是矩阵的大小。

返回值:

返回两个变量的元组( \(w\) , \(v\) )(tuple ( Var , Var )), 其中 \(w\) 是特征值, 维度为(…, \(M\) ), \(v\) 是规范化后的特征向量, 维度为(…, \(M\) , \(M\) )。

代码示例:
>>> x = jt.random((2, 2))
jt.Var([[ 0.9814584  -0.1916754 ]
  [-0.8806686  -0.47373292]], dtype=float32)
>>> jt.linalg.eigh(x)
(jt.Var([-0.54820526  1.4303665 ], dtype=float32), jt.Var([[-0.1916754  -0.9814584 ]
  [-0.47373292  0.8806686 ]], dtype=float32))
jittor.linalg.einsum(string, *args)[源代码]

实现 \(einsum(Einstein Summation)\) 操作。

参数:
  • string(str): 用以描述运算的字符串, 在 \(einsum\) 中指定输入、输出和对决定输出的维度上的操作。

  • args(Sequence[Var]): 待执行 \(einsum\) 操作的输入数据, 可以是一个或多个的数据。

返回值:

经过einsum操作后的Var。

代码示例:
>>> import jittor as jt
>>> a = jt.random((1, 2, 1))
>>> b = jt.random((1, 1, 2))
>>> c = jt.linalg.einsum('bij,bjk->bik', a, b)
>>> print(c)
jt.Var([[[0.10123717 0.6376321 ]
 [0.10257126 0.64603466]]], dtype=float32)
jittor.linalg.inv(x)[源代码]

计算输入矩阵的逆。

参数:
  • x (Var): 输入Var。

返回值:

逆矩阵(Var) 。

代码示例:
>>> x = jt.random((2, 2))
jt.Var([[ 0.9814584  -0.1916754 ]
  [-0.8806686  -0.47373292]], dtype=float32)
>>> jt.linalg.inv(x)
jt.Var([[ 2.520024   -0.4921519 ]
  [-1.1624513  -0.62531066]], dtype=float32)
jittor.linalg.pinv(x)[源代码]

实现对输入矩阵的伪逆运算, 该算子支持梯度反向传播。

参数:
  • x (numpy.ndarray): (…, \(M\) , \(N\) )类型的维度数组或矩阵。

返回值:

返回 \(x\) 的伪逆矩阵(numpy.ndarray) (…, \(N\) , \(M\) );对于每个高阶张量的最后两个维度进行计算, 这两个维度被视为矩阵, 其他高阶维度被视为批次。

代码示例:
>>> import jittor as jt
>>> x = jt.array([[1.0,2.0],[3.0,4.0]])
>>> y = jt.linalg.pinv(x)
>>> print(y)
jt.Var([[-1.9999999  1.       ]
        [ 1.5       -0.5      ]], dtype=float32)
jittor.linalg.qr(x)[源代码]

对输入矩阵进行 \(QR\) 分解。

参数:

x (array): 要进行 \(QR\) 分解的矩阵, 形状为(M,M)。

返回值:

\(QR\) 分解的结果, 是形状为(M,M)的矩阵。

代码示例:
>>> import jittor as jt
>>> x = jt.random((2, 2))
>>> q, r = jt.linalg.qr(x)
>>> print(q, r)
jt.Var([[-0.9639901 -0.2659382]
[-0.2659382  0.9639901]], dtype=float32) jt.Var([[-1.0051305  -1.0211498 ]
[ 0.          0.29402444]], dtype=float32)
jittor.linalg.slogdet(x)[源代码]

计算输入矩阵的行列式的符号与对数值。其采用 LU 分解来计算行列式的符号和自然对数值:

\[det(a) = sign * exp(logdet)\]

其中, sign 表示行列式的符号, logdet 表示行列式的自然对数值。

参数:
  • x (Var): (..., M, M) 维张量, 表示需要求行列式的矩阵。其中, \(M\) 代表矩阵的行列数。

返回值:

返回两个张量, 第一个张量代表行列式的符号, 第二个张量代表行列式的对数值。

代码示例:
>>> import jittor as jt
>>> x = jt.array([[1.0,2.0],[3.0,4.0]])
>>> sign, logdet = jt.linalg.slogdet(x)
>>> print(f"Sign: {sign.data}, LogDet: {logdet.data}")
Sign: [-1.], LogDet: [0.6931472]
jittor.linalg.solve(a, b)[源代码]

求解线性矩阵方程 \(Ax = B\) , 通过以下公式实现:

\[x = A^{-1}B\]

其中 A 是一个非奇异矩阵, B 是一个向量或者矩阵。函数的正向传播将通过Numpy的linalg.solve实现, 反向传播将通过计算梯度公式实现。函数的反向传播在 AB 中存在0的情况下可能会得到不正确的结果。

参数:
  • a(Var): 线性方程组的系数矩阵 \(A\)

  • b(Var): 线性方程组的常数向量 \(B\)

返回值:

线性方程组的解向量(Var)。

代码示例:
>>> a = jt.array([[1., 2.], [3., 4.]]) 
>>> b = jt.array([5., 6.])
>>> from jittor.linalg import solve
>>> print(solve(a, b))
jt.Var([-4.   4.5], dtype=float32)
jittor.linalg.svd(x)[源代码]

计算输入矩阵的奇异值分解。计算方式遵循公式: \(x\) = \(u\) \(s\) \(v\) , 其中 \(u\)\(s\)\(v\) 分别是输入矩阵的左奇异向量、奇异值和右奇异向量。

参数:
  • x (ndarray, jittor.Var): 输入的待分解矩阵。

返回值:

返回奇异值分解后的各部分, u,s,v(ndarray)

代码示例:
>>> from jittor.linalg import svd
>>> X = jt.Var([[1., 0., 0., 0., 2.], [0., 0., 3., 0., 0.], [0., 0., 0., 0., 0.], [0., 2., 0., 0., 0.]])
>>> u, s, v = svd(X)
>>> print(u)
jt.Var([[-0.  1.  0.  0.]
        [ 0.  0.  1.  0.]
        [-1.  0.  0.  0.]
        [ 0.  0.  0.  1.]], dtype=float32)