jit - 是否可以 jit 使用 jax.numpy.unique 的函数?
问题描述
以下代码不起作用:
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()
关于尖位的文档解释说,如果内部数组的形状取决于参数值,则 jit 不起作用。这正是这里的情况。
根据文档,一个潜在的解决方法是指定静态参数。但这不适用于我的情况。几乎每个函数调用的参数都会改变。我已将我的代码拆分为一个预处理步骤,该步骤执行诸如 this之类jnp.unique
的计算,以及一个可以 jitted 的计算步骤。
但是我还是想问一下,是否有一些我不知道的解决方法?
解决方案
不,由于您提到的原因,目前无法jnp.unique
在非静态值上使用。
在类似的情况下,JAX 有时会添加额外的参数,这些参数可用于指定输出的静态大小(例如 中的size
参数jax.numpy.nonzero
),但目前没有为jnp.unique
. 如果这是您想要的,那么值得提交功能请求。
推荐阅读
- objective-c - 字符串文字必须以 '@' 为前缀
- elasticsearch - 在 Elasticsearch 中搜索部分单词的愿望功能不返回任何内容。仅适用于完整的单词
- python - 在 Python 中写入文件时,如何确保数据保存在一行中?
- python - “如果”语句不适用于 np.abs(np.max(f_5.all()))
- javascript - 导航栏问题的下拉菜单
- python - Python 中的未定义行为
- c# - 为什么要在库代码中避免使用 Console.WriteLine?
- django - Django表单不保存
- ios - 如何在 iOS 中将浏览器中的 url 共享到我的应用程序?
- sql - 无法根据 SQL Server 中另一列的条件从列中检索不同的行