jax.experimental.shard_map module

Contents

jax.experimental.shard_map module#

API#

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

Map a function over shards of data.