heat.sparse.dcsx_matrix

Provides DCSR_matrix, a distributed compressed sparse row matrix

Module Contents

class DCSR_matrix(array: torch.Tensor, gnnz: int, gshape: Tuple[int, Ellipsis], dtype: heat.core.types.datatype, split: int | None, device: heat.core.devices.Device, comm: Communication, balanced: bool)[source]

Bases: __DCSX_matrix

Distributed Compressed Sparse Row Matrix. It is composed of PyTorch sparse_csr_tensors local to each process.

Parameters:

  • array (torch.Tensor (layout ==> torch.sparse_csr)) – Local sparse array

  • gnnz (int) – Total number of non-zero elements across all processes

  • gshape (Tuple[int,...]) – The global shape of the array

  • dtype (datatype) – The datatype of the array

  • split (int or None) – If split is not None, it denotes the axis on which the array is divided between processes. DCSR_matrix only supports distribution along axis 0.

  • device (Device) – The device on which the local arrays are using (cpu or gpu)

  • comm (Communication) – The communications object for sending and receiving data

  • balanced (bool or None) – Describes whether the data are evenly distributed across processes.

class DCSC_matrix(array: torch.Tensor, gnnz: int, gshape: Tuple[int, Ellipsis], dtype: heat.core.types.datatype, split: int | None, device: heat.core.devices.Device, comm: Communication, balanced: bool)[source]

Bases: __DCSX_matrix

Distributed Compressed Sparse Column Matrix. It is composed of PyTorch sparse_csc_tensors local to each process.

Parameters:

  • array (torch.Tensor (layout ==> torch.sparse_csc)) – Local sparse array

  • gnnz (int) – Total number of non-zero elements across all processes

  • gshape (Tuple[int,...]) – The global shape of the array

  • dtype (datatype) – The datatype of the array

  • split (int or None) – If split is not None, it denotes the axis on which the array is divided between processes. DCSR_matrix only supports distribution along axis 0.

  • device (Device) – The device on which the local arrays are using (cpu or gpu)

  • comm (Communication) – The communications object for sending and receiving data

  • balanced (bool or None) – Describes whether the data are evenly distributed across processes.