pytc.vmc.sharding.create_mesh¶

pytc.vmc.sharding.create_mesh(devices=None, axis_name='walkers')[source]¶

Create a 1-D device mesh for walker-parallel sharding.

Parameters:

devices (sequence of jax.Device, optional) – Devices to include. Defaults to jax.devices().
axis_name (str) – Name for the walker-parallel axis (default "walkers").

Return type:

jax.sharding.Mesh