pytc.vmc.sharding.replicate

pytc.vmc.sharding.replicate(pytree, mesh)[source]

Replicate a pytree on every device.

Parameters:

mesh (jax.sharding.Mesh)