pytc.vmc.sharding._slice_along_first_axis¶ pytc.vmc.sharding._slice_along_first_axis(pytree, start, end)[source]¶ Slice a pytree along the leading axis for all non-scalar leaves. Parameters: start (int) end (int)