mlx.core.linalg.lu

目录

mlx.core.linalg.lu#

lu(a: array, *, stream: None | Stream | Device = None) → Tuple[array, array, array]#

计算给定矩阵 A 的 LU 分解。

注意，与 scipy.linalg.lu 的默认行为不同，这里的枢轴是索引。要重构输入，对于 2 维使用 L[P, :] @ U，对于超过 2 维使用 mx.take_along_axis(L, P[..., None], axis=-2) @ U。

要构建完整的置换矩阵，请执行以下操作：

P = mx.put_along_axis(mx.zeros_like(L), p[..., None], mx.array(1.0), axis=-1)
参数:

  • a (array) – 输入数组。

  • stream (Stream, optional) – 流或设备。默认为 None，此时使用默认设备的默认流。

返回:

数组 p、L 和 U，满足 A = L[P, :] @ U。

返回类型:

tuple(array, array, array)

目录