pytc.vmc.sharding.sharded_batched_vmap¶

pytc.vmc.sharding.sharded_batched_vmap(fn, max_batch_size, mesh=None, in_axes=0, out_axes=0)[source]¶

A version of folx.batched_vmap that works across multiple devices.

It uses jax.shard_map to distribute the computation across devices, and then folx.batched_vmap on each device to process the local shard in memory-efficient chunks.