jax.custom_batching.sequential_vmap

jax.custom_batching.sequential_vmap(f)

是 custom_vmap 的一个特殊情况，它使用循环实现。

用 sequential_vmap 装饰的函数在批处理时将在循环中按顺序调用。这对于本身不支持批处理维度的函数很有用。

例如

>>> @jax.custom_batching.sequential_vmap
... def f(x):
...   jax.debug.print("{}", x)
...   return x + 1
...
>>> jax.vmap(f)(jnp.arange(3))
0
1
2
Array([1, 2, 3], dtype=int32)

其中打印语句表明此 vmap() 是使用循环生成的。

有关 custom_vmap 的更多详细信息，请参阅文档。