Python Jax中的vmap ops.index_更新

我有下面的代码，它使用了一个简单的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_轴必须是int、None或（嵌套）容器 用那些类型的叶子，但是 追踪

我有点困惑，因为范围是int类型的，所以我不太确定发生了什么

最后，我将尽可能优化这个小部件，以获得最短的时间。

可以表达单个操作跨输入的多个轴独立应用的功能。您的函数有点不同：您对单个输入迭代应用了单个操作

幸运的是，JAX提供了可以处理这种情况的工具。实现将如下所示：

来自jax导入lax的

def参数更新（y，ind）：
返回jax.ops.index_update（y，jax.ops.index[：，ind]，jscp.convalve（pulse_20，y[：，ind]）[：-19]），ind
@jax.jit
def过滤器_jax2（y）：
ranger=jnp.arange（len（y））
返回松弛扫描（paraUpdate，y，ranger）[0]
打印（np.allclose（filter\ujax（jnpData）、filter\ujax2（jnpData）））
#真的
%timeit filter_jax（jnpData）.block_直到_就绪（）
#10个回路，最佳3个：每个回路28.6毫秒
%timeit filter_jax2（jnpData）.block_直到_就绪（）
#1000个回路，最佳3个：每个回路519µs

如果更改算法以便将操作应用于数组中的每一列，而不是前N列，则可以使用

vmap

表示如下：

@jax.jit
def过滤器_jax3（y）：
f=λ列：jscp.卷积（脉冲_20，列）[：-19]
返回jax.vmap（f，in_轴=1，out_轴=1）（y）

这解决了我遇到的问题，谢谢！不过，我确实有另一个问题，我注意到JAX等价物不再匹配此代码的原始纯python版本，这是什么原因造成的

def filter_cpu（y）：y[：，range（0，len（y））]=signal.convalve（pulse_20[np.newaxis，：]，y[：，range（0，len（y））]）[：，：-19]返回y

JAX默认为32位计算，除非设置X64标志；scipy默认为64位计算。这可能就是不同之处。还请注意，如果您想直接使用jax，它有一个2D卷积：