gpt4 book ai didi

python - 如何将 jax vmap 用于嵌套循环?

转载 作者:行者123 更新时间:2023-12-04 14:48:52 27 4
gpt4 key购买 nike

我想使用 vmap 来矢量化此代码以提高性能。

def matrix(dataA, dataB):
return jnp.array([[func(a, b) for b in dataB] for a in dataA])
matrix(data, data)
我试过这个:
def f(x, y):
return func(x, y)
mapped = jax.vmap(f)
mapped(data, data)
但这仅给出对角线条目。
基本上我有一个向量 data = [1,2,3,4,5] (示例),我想得到一个矩阵,使得每个条目 (i, j)矩阵是 f(data[i], data[j]) .因此,生成的矩阵形状将为 (len(data), len(data)) .

最佳答案

jax.vmap一次映射一组轴。如果要映射两个独立的轴集,可以通过嵌套两个 vmap 来实现。转换:

mapped = jax.vmap(jax.vmap(f, in_axes=(None, 0)), in_axes=(0, None))
result = mapped(data, data)

关于python - 如何将 jax vmap 用于嵌套循环?,我们在Stack Overflow上找到一个类似的问题: https://stackoverflow.com/questions/69429846/

27 4 0
Copyright 2021 - 2024 cfsdn All Rights Reserved 蜀ICP备2022000587号
广告合作:1813099741@qq.com 6ren.com