并发

并发#

JAX 对 Python 并发支持有限。

客户端可以从不同的 Python 线程并发调用 JAX API(例如 jit()grad())。

不允许从多个线程并发操作 JAX 跟踪值。换句话说,虽然允许从多个线程调用使用 JAX 跟踪的函数(例如 jit()),但您不能使用线程在传递给 jit() 的函数 f 的实现中操作 JAX 值。如果您这样做,最可能的结果是 JAX 会出现一个神秘的错误。