pytc.vmc.sharding.shard_vmap¶ pytc.vmc.sharding.shard_vmap(fn, mesh=None, in_axes=0, out_axes=0)[source]¶ A version of jax.vmap that works across multiple devices using shard_map.