问题描述
以下代码不起作用:
def get_unique(arr):
return jnp.unique(arr)
get_unique = jit(get_unique)
get_unique(jnp.ones((10,)))
错误消息是关于 jnp.unique
的使用:
FilteredStackTrace: jax._src.errors.ConcretizationTypeError: Abstract tracer value encountered where concrete value is expected: Traced<ShapedArray(float32[10])>with<DynamicJaxprTrace(level=0/1)>
The error arose in jnp.unique()
documentation on sharp bits 解释了如果内部数组的形状取决于参数值,则 jit 不起作用。这正是这里的情况。
根据文档,一个潜在的解决方法是指定静态参数。但这不适用于我的情况。几乎每个函数调用的参数都会改变。我已将我的代码拆分为一个预处理步骤,该步骤执行诸如 jnp.unique
之类的计算,以及一个可以 jitted 的计算步骤。
但我还是想问一下,是否有一些我不知道的解决方法?
解决方法
不,由于您提到的原因,目前无法对非静态值使用 jnp.unique
。
在类似的情况下,JAX 有时会添加额外的参数,用于指定输出的静态大小(例如,jax.numpy.nonzero
中的 size
参数),但目前没有类似的实现jnp.unique
。如果这是您想要的,那么提交 feature request 是值得的。