pylibwholegraph.torch.create_wholememory_tensor

pylibwholegraph.torch.create_wholememory_tensor(

comm: WholeMemoryCommunicator,

memory_type: str,

memory_location: str,

sizes: List[int],

dtype: dtype,

strides: List[int],

tensor_entry_partition: List[int] | None = None,

)

Create empty WholeMemory Tensor. Now only support dim = 1 or 2 :param comm: WholeMemoryCommunicator :param memory_type: WholeMemory type, should be continuous, chunked or distributed :param memory_location: WholeMemory location, should be cpu or cuda :param sizes: size of the tensor :param dtype: data type of the tensor :param strides: strides of the tensor :param tensor_entry_partition: rank partition based on entry;

tensor_entry_partition[i] determines the entry count of rank i and shoud be a positive integer; the sum of tensor_entry_partition should equal to total entry count; entries will be equally partitioned if None

Returns:

Allocated WholeMemoryTensor