pytc.vmc.sharding.shard_map_wrap¶

pytc.vmc.sharding.shard_map_wrap(fn, mesh, in_specs, out_specs)[source]¶

Apply shard_map with compatibility flags across JAX versions.