jax.experimental.sparse.linalg.lobpcg_standard#

jax.experimental.sparse.linalg.lobpcg_standard(A, X, m=100, tol=None)[源代码]#

使用 LOBPCG 例程计算 top-k 标准特征值。

LOBPCG [1] 代表局部最优块预处理共轭梯度法。该方法能够以加速器友好的方式找到 top-k 特征向量。

这个最初的实验版本有一些注意事项。

  • 仅支持标准特征值问题 A U = lambda U,不支持一般特征值问题。

  • 梯度代码不可用。

  • f64 仅在 jnp.linalg.eigh 支持该类型时才有效。

  • 尚不支持找到最小的特征向量。因此,我们还不支持预处理,这主要是在这种情况下需要的。

该实现基于 [2] 和 [3];但是,我们以几种方式偏离这些来源,以提高稳健性或促进实施

  • 尽管迭代成本增加,但我们始终保持块搜索方向的正交基。

  • 我们更改了收敛标准;请参阅 tol 参数。

  • 软锁定 [4] 是故意未实现的;它依赖于选择适当的特定于问题的容差,以防止接近收敛时由于接近 0 残差的灾难性消除而导致的爆炸。相反,所实现的方法倾向于截断迭代基础。

[1]: http://ccm.ucdenver.edu/reports/rep149.pdf [2]: https://arxiv.org/abs/1704.07458 [3]: https://arxiv.org/abs/0705.2626 [4]: DOI 10.13140/RG.2.2.11794.48327

参数:
  • A (jax.Array | Callable[[jax.Array], jax.Array]) – 一个 (n, n) 数组,表示一个 Hermitian 方阵或具有其动作的可调用对象。

  • X (jax.Array) – 一个 (n, k) 数组,表示 k 个所需 top 特征向量的初始搜索方向。这不需要是正交的,但必须在数值上线性独立(X 将被正交化)。请注意,我们必须有 0 < k * 5 < n

  • m (int) – 最大整数迭代计数;LOBPCG 将只探索(Krylov 基的子空间){X, A X, A^2 X, …, A^m X}

  • tol (jax.Array | float | None) – 浮点收敛容差;当特征对 (lambda, v) 的残差 L2 范数 r = |A v - lambda v| 低于 tol * 10 * n * (lambda + |A v|) 时,该特征对收敛,这粗略估计了理想特征向量的最坏情况浮点误差。如果所有 k 个特征向量都满足容差比较,则 LOBPCG 提前退出。如果保留为 None,则将其设置为 A.dtype 的浮点 epsilon。

返回:

theta, U, i,其中 theta 是一个 (k,) 特征值数组,U 是一个 (n, k) 特征向量数组,i 是执行的迭代次数。

引发:

ValueError – 如果 A,X 的 dtypes 或 n 的维度不匹配,或者 k 太大(仅支持 k * 5 < n),或者 k == 0