heat.sparse.dcsr_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)

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.

global_indptr() → heat.core.dndarray.DNDarray

Global indptr of the DCSR_matrix as a DNDarray

is_distributed() → bool

Determines whether the data of this DCSR_matrix is distributed across multiple processes.

counts_displs_nnz() → Tuple[Tuple[int], Tuple[int]]

Returns actual counts (number of non-zero items per process) and displacements (offsets) of the DCSR_matrix. Does not assume load balance.

astype(dtype, copy=True) → DCSR_matrix

Returns a casted version of this matrix. Casted matrix is a new matrix of the same shape but with given type of this matrix. If copy is True, the same matrix is returned instead.

Parameters:

  • dtype (datatype) – HeAT type to which the matrix is cast

  • copy (bool, optional) – By default the operation returns a copy of this matrix. If copy is set to False the cast is performed in-place and this matrix is returned

__repr__() → str

Computes a printable representation of the passed DCSR_matrix.