def laid_out_slice_num(self, tensor_shape):
    """A LaidOutTensor with an int32 scalar, identical for identical slices.

    This is useful for synchronizing random operations.

    Args:
      tensor_shape: a TensorShape
    Returns:
      a LaidOutTensor where each slice is an integer scalar.
    """
    ret = self.slicewise(lambda: tf.to_int32(0))
    tensor_layout = self.tensor_layout(tensor_shape)
    for mesh_axis in tensor_layout.tensor_axis_to_mesh_axis:
      if mesh_axis is not None:
        def my_fn(x, pcoord, mesh_dim_size):
          return x * mesh_dim_size + pcoord
        ret = self.slicewise(
            my_fn, ret, self.laid_out_pcoord(mesh_axis),
            self.shape[mesh_axis].size)
    return ret
def laid_out_slice_num ( self , tensor_shape ) : ret = self . slicewise ( lambda : tf . to_int32 ( 0 ) ) tensor_layout = self . tensor_layout ( tensor_shape ) for mesh_axis in tensor_layout . tensor_axis_to_mesh_axis : if mesh_axis is not None : def my_fn ( x , pcoord , mesh_dim_size ) : return x * mesh_dim_size + pcoord ret = self . slicewise ( my_fn , ret , self . laid_out_pcoord ( mesh_axis ) , self . shape [ mesh_axis ] . size ) return ret
A LaidOutTensor with an int32 scalar, identical for identical slices.

    This is useful for synchronizing random operations.

    Args:
      tensor_shape: a TensorShape
    Returns:
      a LaidOutTensor where each slice is an integer scalar.