我怎样才能通过 JAX 加速这个循环代码？寻找碰撞的最近邻居

Question

我正在尝试在另一个 SO question 上使用 JAX 来评估代码的 JAX 适用性和性能（有关于代码功能的有用信息）。为此，我通过 jax.numpy (jnp) 等效方法修改了代码（由于我对 JAX 的经验很少，用等效的 jnp 代码替换 NumPy 相关代码并不像我想象的那么容易，也许它可以写得更好）。最后，我用前代码 (optimized algorithm) 检查了结果，结果是一样的，但是 JAX 需要 7.5 秒，而前一个示例需要 0.10 秒（使用 Colab ）。我认为这么长的运行时间可能与代码中的for 循环有关，它可能被 JAX 相关模块所​​取代，例如fori-loop 或 vectorization 和……； 但我不知道必须进行哪些更改以及如何更改才能使这段代码在性能和速度方面令人满意（使用 JAX）。

import numpy as np
from scipy.spatial import cKDTree, distance
import jax
from jax import numpy as jnp
jax.config.update("jax_enable_x64", True)


# ---------------------------- input data ----------------------------
""" For testing by prepared files:
radii = np.load('a.npy')
poss = np.load('b.npy')
"""

rnd = np.random.RandomState(70)
data_volume = 1000

radii = rnd.uniform(0.0005, 0.122, data_volume)
dia_max = 2 * radii.max()

x = rnd.uniform(-1.02, 1.02, (data_volume, 1))
y = rnd.uniform(-3.52, 3.52, (data_volume, 1))
z = rnd.uniform(-1.02, -0.575, (data_volume, 1))
poss = np.hstack((x, y, z))
# --------------------------------------------------------------------


# @jax.jit
def ends_gap(poss, dia_max):
    particle_corsp_overlaps = jnp.array([], dtype=np.float64)

    # kdtree = cKDTree(poss)                                                                                              # Using SciPy

    for particle_idx in range(len(poss)):

        cur_point = poss[particle_idx]
        # nears_i_ind = jnp.array(kdtree.query_ball_point(cur_point, r=dia_max, return_sorted=True), dtype=np.int64)      # Using SciPy
        
        # Using NumPy
        unshared_idx = jnp.delete(jnp.arange(len(poss)), particle_idx)
        poss_without = poss[unshared_idx]
        dist_max = radii[particle_idx] + radii.max()

        lx_limit_idx = poss_without[:, 0] <= poss[particle_idx][0] + dist_max
        ux_limit_idx = poss_without[:, 0] >= poss[particle_idx][0] - dist_max
        ly_limit_idx = poss_without[:, 1] <= poss[particle_idx][1] + dist_max
        uy_limit_idx = poss_without[:, 1] >= poss[particle_idx][1] - dist_max
        lz_limit_idx = poss_without[:, 2] <= poss[particle_idx][2] + dist_max
        uz_limit_idx = poss_without[:, 2] >= poss[particle_idx][2] - dist_max

        nears_i_ind = jnp.where(lx_limit_idx & ux_limit_idx & ly_limit_idx & uy_limit_idx & lz_limit_idx & uz_limit_idx)[0]

        # assert len(nears_i_ind) > 0

        # if len(nears_i_ind) <= 1:
        #     continue

        nears_i_ind = nears_i_ind[nears_i_ind != particle_idx]

        # dist_i = distance.cdist(poss[tuple(nears_i_ind[None, :])], cur_point[None, :]).squeeze()                        # Using SciPy
        dist_i = jnp.linalg.norm(poss[tuple(nears_i_ind[None, :])] - cur_point[None, :], axis=-1)                     # Using NumPy
        contact_check = dist_i - (radii[tuple(nears_i_ind[None, :])] + radii[particle_idx])

        connected = contact_check[contact_check <= 0]
        particle_corsp_overlaps = jnp.concatenate((particle_corsp_overlaps, connected))

        contacts_ind = jnp.where(contact_check <= 0)[0]
        contacts_sec_ind = jnp.array(nears_i_ind)[contacts_ind]
        sphere_olps_ind = jnp.sort(contacts_sec_ind)

        ends_ind_mod_temp = jnp.array([jnp.repeat(particle_idx, len(sphere_olps_ind)), sphere_olps_ind], dtype=np.int64).T
        if particle_idx > 0:   # ---> these 4-lines perhaps be better to be substituted by just one-line list appending as "ends_ind.append(ends_ind_mod_temp)"
            ends_ind = jnp.concatenate((ends_ind, ends_ind_mod_temp))
        else:
            ends_ind = jnp.array(ends_ind_mod_temp, dtype=np.int64)

    ends_ind_org = ends_ind
    ends_ind, ends_ind_idx = jnp.unique(jnp.sort(ends_ind_org), axis=0, return_index=True)
    gap = jnp.array(particle_corsp_overlaps)[ends_ind_idx]

    return gap, ends_ind, ends_ind_idx, ends_ind_org

我曾尝试在此代码上使用@jax.jit，但它显示错误：TracerArrayConversionError 或ConcretizationTypeError on COLAB TPU：

使用 SciPy：

TracerArrayConversionError：numpy.ndarray 转换方法 array() 在 JAX Tracer 对象上被调用 Tracedwith 在跟踪函数 ends_gap 时 :1 对于 jit，这个具体值不是 在 Python 中可用，因为它取决于参数的值 '可能'。看 https://jax.readthedocs.io/en/latest/errors.html#jax.errors.TracerArrayConversionError

使用 NumPy：

ConcretizationTypeError：在哪里遇到抽象跟踪器值 预期的具体价值： Tracedwith jnp.nonzero 的 size 参数必须静态指定才能使用 JAX 转换中的 jnp.nonzero。跟踪函数时 对于 jit，ends_gap at :1，这个 具体值在 Python 中不可用，因为它取决于 参数 'poss' 和 'dia_max' 的值。

看 https://jax.readthedocs.io/en/latest/errors.html#jax.errors.ConcretizationTypeError

如果可以通过使用 JAX 传递这些问题（如果可能，jax.jit）来加速此代码，我将不胜感激。如何利用 JAX 在 CPU 和 GPU 或 TPU 上获得最佳性能？

---

准备样品测试数据：
a.npy = Radii data
b.npy = Poss data

更新

---

这个问题的主要目的是如何修改代码以获得使用 JAX 库的最佳性能

我已在基于jakevdp answer 的代码上注释了 SciPy 相关行，并取消注释等效的 NumPy 相关部分。
为了获得更好的答案，我对一些重要的主题进行了编号：

1. scikit-learn BallTree相关方法是否兼容JAX？？在内存使用方面（对于可能的矢量化），这种方法可以成为 SciPy cKDTree 的一个很好的替代方案。
2. 如何最好地处理代码中的循环部分，使用fori_loop 或将循环的代码行放入函数中，然后向量化、jitting 或...？？

• 我在准备使用 fori_loop 的代码时遇到问题。使用fori_loop 所做的工作可以从下面的代码行中理解，其中particle_corsp_overlaps 是定义函数的输入（这个函数只包含循环部分）。如果推荐使用fori_loop，那么展示如何做到这一点将很有用。

particle_corsp_overlaps, ends_ind = jax.lax.fori_loop(0, len(poss), jax_loop, particle_corsp_overlaps)

• 我将 NumPy 部分放在一个函数中，由@jax.jit 进行 jitting 以检查其提高性能的能力（我不知道它有多大帮助）。它收到与poss 相关的错误ConcretizationTypeError（--> 形状取决于跟踪值）。因此，我尝试通过从functools 导入partial 来使用@partial(jax.jit, static_argnums=0) 装饰器，但现在我收到以下错误；如果推荐这种方式如何解决它，例如为：

@partial(jax.jit, static_argnums=0)
def ends_gap(poss):

    for particle_idx in range(len(poss)):

        cur_point = poss[particle_idx]

        unshared_idx = jnp.delete(jnp.arange(len(poss)), particle_idx)
        poss_without = poss[unshared_idx]
        dist_max = radii[particle_idx] + radii.max()

        lx_limit_idx = poss_without[:, 0] <= poss[particle_idx][0] + dist_max
        ux_limit_idx = poss_without[:, 0] >= poss[particle_idx][0] - dist_max
        ly_limit_idx = poss_without[:, 1] <= poss[particle_idx][1] + dist_max
        uy_limit_idx = poss_without[:, 1] >= poss[particle_idx][1] - dist_max
        lz_limit_idx = poss_without[:, 2] <= poss[particle_idx][2] + dist_max
        uz_limit_idx = poss_without[:, 2] >= poss[particle_idx][2] - dist_max

        nears_i_ind = jnp.where(lx_limit_idx & ux_limit_idx & ly_limit_idx & uy_limit_idx & lz_limit_idx & uz_limit_idx)[0]
        nears_i_ind = nears_i_ind[nears_i_ind != particle_idx]

        dist_i = jnp.linalg.norm(poss[tuple(nears_i_ind[None, :])] - cur_point[None, :], axis=-1) 

ValueError：不支持不可散列的静态参数。一个错误 在尝试散列时调用“nearest_neighbors_jax”期间发生 类型的对象，[[ 8.42519143e-01 1.37693422e+00 -7.97775882e-01] [-3.31436445e-01 -1.67346250e+00 -8.61069684e-01] [-1.57500126e-01 -1.17502591e+00007].8791e+00007。错误是：TypeError: unhashable type: 'DeviceArray'

由于卡在这个简短定义的函数中，我没有将整个循环体放入函数中。如果可能的话，创建一个包含所有循环体的函数（可以是 jit 或 ...）是很有趣的。

1. 是否可以使用 jax 方法将 4 行 ends_ind 相关 if-else 语句写在一行中，以避免 if 在 jitting 或...期间可能出现问题？

Answer 1

JAX 不能用于优化一般的 numpy/scipy 代码，但它可以用于优化/编译用 JAX 编写的代码。

您的示例围绕 scipy 的 cKDTree 的使用展开。这在 JAX 中没有实现，因此无法在 JAX 中进行优化或编译，并且在 jitted 函数中使用它会导致您看到的错误。如果你想在 JAX 中使用 KD 树，你必须找到一个用 JAX 实现的树。我不知道任何这样的代码。

至于为什么在这里将np 替换为jnp 时代码会变慢，这是因为您实际上只是将JAX 用作备用数组容器。每次将 JAX 数组传递给 cKDTree 调用时，都必须将其转换为 numpy 数组，然后必须将结果转换回 JAX 数组。这种额外的数据移动会增加每次调用的开销，从而使结果变慢。这并不是因为 JAX 本身很慢，而是因为您并没有真正将 JAX 用作任何东西，而是在将数据转换回 numpy 之前临时存储数据的一种方式。

一般可以通过将函数包装在jax.jit中来减少这种开销，但是如前所述，这与scipy.spatial.cKDTree这样的非jax代码不兼容。

我怀疑您最好的做法是避免使用 JAX，而只使用 numpy/scipy 和 cKDTree。我不知道任何与 JAX 兼容的基于树的邻居搜索算法的实现，对于大型数组，完整的蛮力方法不会与 cKDTree 竞争。