Taichi 分布式

Taichi 是否可以结合Ray或者Dask等分布式计算的框架?

尝试了一下在Ray里套皮貌似寄了~

import ray
import taichi as ti

# ti.init(ti.gpu)

@ti.func
def is_prime(n: int):
    result = True
    for k in range(2, int(n ** 0.5) + 1):
        if n % k == 0:
            result = False
            break
    return result

@ti.kernel
def count_primes(l: int, r: int) -> int:
    count = 0
    for k in range(l, r):
        if is_prime(k):
            count += 1

    return count

@ray.remote
def count_wrapper(l:int, r:int) -> int:
    return count_primes(l, r)

ray.init()
futures = []
for i in range(2):
    worker = count_wrapper.remote(100*i, 100*i+100)
    futures.append(worker)

res = ray.get(futures)
print(res)