jax.experimental.shard_map 模块

内容

jax.experimental.shard_map 模块#

API#

shard_map(f, mesh, in_specs, out_specs[, ...])

将函数映射到数据的分片上。