我有下面的代码,它使用了一个简单的for循环。我只是想知道是否有一种方法可以将其vmap?以下是原始代码:
import numpy as np
import jax.numpy as jnp
import jax.scipy.signal as jscp
from scipy import signal
import jax
data = np.random.rand(192,334)
a = [1,-1.086740193996892,0.649914553946275,-0.124948974636730]
b = [0.054778173164082,0.164334519492245,0.164334519492245,0.054778173164082]
impulse = signal.lfilter(b, a, [1] + [0]*99)
impulse_20 = impulse[:20]
impulse_20 = jnp.asarray(impulse_20)
@jax.jit
def filter_jax(y):
for ind in range(0, len(y)):
y = jax.ops.index_update(y, jax.ops.index[:, ind], jscp.convolve(impulse_20, y[:,ind])[:-19])
return y
jnpData = jnp.asarray(data)
%timeit filter_jax(jnpData).block_until_ready()
下面是我使用vmap的尝试:
def paraUpdate(y, ind):
return jax.ops.index_update(y, jax.ops.index[:, ind], jscp.convolve(impulse_20, y[:,ind])[:-19])
@jax.jit
def filter_jax2(y):
ranger = range(0, len(y))
return jax.vmap(paraUpdate, y)(ranger)
但我收到以下错误:
TypeError: vmap in_axes must be an int, None, or (nested) container with those types as leaves, but got Traced<ShapedArray(float32[192,334])>with<DynamicJaxprTrace(level=0/1)>.
我有点困惑,因为范围是int类型的,所以我不太确定发生了什么
最后,我将尽可能地优化这个小部件,以获得最短的时间
^{} 可以表示单个操作在输入的多个轴上独立应用的功能。您的函数有点不同:您对单个输入迭代应用了单个操作
幸运的是,JAX提供了^{} ,可以处理这种情况。实现将如下所示:
如果更改算法,以便将操作应用于数组中的每个列,而不是第一个N列,则可以使用以下
vmap
表示:相关问题 更多 >
编程相关推荐