pytc.vmc.sharding.get_vmap_fn¶

pytc.vmc.sharding.get_vmap_fn(max_vmap_batch_size=0, mesh=None)[source]¶

Return the appropriate vmap implementation.

Automatically detects if multiple GPUs are available. If so, and max_vmap_batch_size > 0, uses sharded_batched_vmap which combines sharding and batching. If max_vmap_batch_size == 0, uses shard_vmap for sharded data-parallel execution.

Returns:

A vmap-like function with (fn, in_axes, ...) signature.

Return type:

callable

Parameters:

max_vmap_batch_size (int)
mesh (jax.sharding.Mesh | None)