pygho.backend package

Submodules

pygho.backend.MaTensor module

class pygho.backend.MaTensor.MaskedTensor(data: Tensor, mask: BoolTensor, padvalue: float = 0.0, is_filled: bool = False)[source]

Bases: object

Represents a masked tensor with optional padding values. This class allows you to work with tensors that have a mask indicating valid and invalid values. You can perform various operations on the masked tensor, such as filling masked values, computing sums, means, maximums, minimums, and more.

Parameters:

  • data (Tensor): The underlying data tensor of shape (*maskedshape, *denseshape)

  • mask (BoolTensor): The mask tensor of shape (*maskedshape) where True represents valid values, and False` represents invalid values.

  • padvalue (float, optional): The value to use for padding. Defaults to 0.

  • is_filled (bool, optional): Indicates whether the invalid values have already been filled to the padvalue. Defaults to False.

Attributes:

  • data (Tensor): The underlying data tensor.

  • mask (BoolTensor): The mask tensor.

  • fullmask (BoolTensor): The mask tensor after broadcasting to match the data’s dimensions.

  • padvalue (float): The padding value.

  • shape (torch.Size): The shape of the data tensor.

  • masked_dim (int): The number of dimensions in maskedshape.

  • dense_dim (int): The number of dimensions in denseshape.

  • maskedshape (torch.Size): The shape of the tensor up to the masked dimensions.

  • denseshape (torch.Size): The shape of the tensor after the masked dimensions.

Methods:

  • fill_masked_(self, val: float = 0) -> None: In-place fill of masked values.

  • fill_masked(self, val: float = 0) -> Tensor: Return a tensor with masked values filled with the specified value.

  • to(self, device: torch.DeviceObjType, non_blocking: bool = True): Move the tensor to the specified device.

  • sum(self, dims: Union[Iterable[int], int], keepdim: bool = False): Compute the sum of masked values along specified dimensions.

  • mean(self, dims: Union[Iterable[int], int], keepdim: bool = False): Compute the mean of masked values along specified dimensions.

  • max(self, dims: Union[Iterable[int], int], keepdim: bool = False): Compute the maximum of masked values along specified dimensions.

  • min(self, dims: Union[Iterable[int], int], keepdim: bool = False): Compute the minimum of masked values along specified dimensions.

  • diag(self, dims: Iterable[int]): Extract diagonals from the tensor. The dimensions in dims will be take diagonal and put at dims[0]

  • unpooling(self, dims: Union[int, Iterable[int]], tarX): Perform unpooling operation along specified dimensions.

  • tuplewiseapply(self, func: Callable[[Tensor], Tensor]): Apply a function to each element of the masked tensor.

  • diagonalapply(self, func: Callable[[Tensor, LongTensor], Tensor]): Apply a function to diagonal elements of the masked tensor.

  • add(self, tarX, samesparse: bool): Add two masked tensors together.

  • catvalue(self, tarX, samesparse: bool): Concatenate values of two masked tensors.

add(tarX, samesparse: bool)[source]
catvalue(tarX, samesparse: bool)[source]
property data: Tensor
property dense_dim
property denseshape
diag(dims: Iterable[int])[source]

put the reduced output to dim[0]

diagonalapply(func: Callable[[Tensor, LongTensor], Tensor])[source]
fill_masked(val: float = 0) → Tensor[source]

return a tensor with masked values filled with val.

fill_masked_(val: float = 0) → None[source]

inplace fill the masked values

property fullmask: BoolTensor
property mask: BoolTensor
property masked_dim
property maskedshape
max(dims: Iterable[int] | int, keepdim: bool = False)[source]
mean(dims: Iterable[int] | int, keepdim: bool = False)[source]
min(dims: Iterable[int] | int, keepdim: bool = False)[source]
property padvalue: float
property shape: Size
sum(dims: Iterable[int] | int, keepdim: bool = False)[source]
to(device: DeviceObjType, non_blocking: bool = True)[source]

move data to some device

tuplewiseapply(func: Callable[[Tensor], Tensor])[source]
unpooling(dims: int | Iterable[int], tarX)[source]
pygho.backend.MaTensor.filterinf(X: Tensor, filled_value: float = 0)[source]

Replaces positive and negative infinity values in a tensor with a specified value.

Args:

  • X (Tensor): The input tensor.

  • filled_value (float, optional): The value to replace positive and negative infinity values with (default: 0).

Returns:

  • Tensor: A tensor with positive and negative infinity values replaced by the specified filled_value.

Example:

input_tensor = torch.tensor([1.0, 2.0, torch.inf, -torch.inf, 3.0])
result = filterinf(input_tensor, filled_value=999.0)

pygho.backend.Mamamm module

pygho.backend.Mamamm.batched_tensordot(A: Tensor, catdim1: int, dim1: int, B: Tensor, catdim2: int, dim2: int) → Tensor[source]

Perform a batched tensordot matrix operation.

This function computes the tensordot product of two tensors A and B, where A and B are batched tensors with specified concatenation dimensions catdim1 and catdim2, and contraction dimensions dim1 and dim2.

Args:

  • A (Tensor): The first batched tensor of shape (catshape1, broadcastshape).

  • catdim1 (int): The length of catshape1.

  • dim1 (int): The contraction dimension along catdim1 of the first tensor.

  • B (Tensor): The second batched tensor of shape (catshape2, broadcastshape)..

  • catdim2 (int): The length of catshape2.

  • dim2 (int): The contraction dimension along catdim2 of the second tensor.

Returns:

  • Tensor: The result of the batched tensordot operation of shape (*catshape1dim1, *catshape2dim2, *broadcastshape), where densedim is the common dense dimension of A and B.

Notes:

  • catdim1 and catdim2 specify the number of concatenation dimensions of A and B, respectively.

  • dim1 and dim2 specify the contraction dimensions along catdim1 and catdim2, respectively.

  • The function uses optimized paths for specific cases (e.g., when catdim1=2 and catdim2=2).

pygho.backend.Mamamm.broadcast_denseshape(A: Tensor, densedim1: int, B: Tensor, densedim2: int) → Tuple[Tensor, Tensor][source]

This function broadcasts the dense shapes of tensors A and B to the same by adding dimensions of size 1.

Args:

  • A (Tensor): The first tensor.

  • densedim1 (int): The number of dense dimension of the first tensor.

  • B (Tensor): The second tensor.

  • densedim2 (int): The number of dense dimension of the second tensor.

Returns:

  • Tuple[Tensor, Tensor]: A tuple containing the broadcasted tensors A and B with compatible dense shapes.

Notes:

  • This function adds dimensions with size 1 to the smaller dense shape until both dense shapes match.

pygho.backend.Mamamm.mamamm(A: MaskedTensor, dim1: int, B: MaskedTensor, dim2: int, mask: BoolTensor, broadcast_firstdim: bool = True) → MaskedTensor[source]

Batched masked matrix multiplication of two MaskedTensors.

This function performs batched matrix multiplication between two MaskedTensors A and B, where the masked dimensions dim1 and dim2 are contracted. The result is a new MaskedTensor with the specified mask.

Args:

  • A (MaskedTensor): The first MaskedTensor with shape (B,* maskedshape1,*denseshapeshape).

  • dim1 (int): The masked dimension to contract in the first tensor A.

  • B (MaskedTensor): The second MaskedTensor with shape (B,* maskedshape2,*denseshapeshape).

  • dim2 (int): The masked dimension to contract in the second tensor B.

  • mask (BoolTensor): The mask to apply to the resulting MaskedTensor.

  • broadcast_firstdim (bool, optional): If True, broadcast the first dimension (batch dimension) of A and B to ensure compatibility. Default is True.

Returns:

  • MaskedTensor: A new MaskedTensor with shape (B,* maskedshape1dim1,* maskedshape2dim2,*denseshapeshape) and the specified mask.

Notes:

  • This function performs batched matrix multiplication between two MaskedTensors, contracting the specified masked dimensions.

pygho.backend.SpTensor module

class pygho.backend.SpTensor.SparseTensor(indices: LongTensor, values: Tensor | None = None, shape: List[int] | None = None, is_coalesced: bool = False, reduce: str = 'sum')[source]

Bases: object

Represents a sparse tensor in coo format.

This class allows you to work with sparse tensors represented by indices and values. It provides various operations such as sum, max, mean, unpooling, diagonal extraction, and more.

Parameters: - indices (LongTensor): The indices of the sparse tensor, of shape (#sparsedim, #nnz). - values (Optional[Tensor]): The values associated with the indices, of shape (#nnz,*denseshapeshape). Should have the same number of nnz as indices. Defaults to None. - shape (Optional[List[int]]): The shape of the sparse tensor. If None, it is computed from the indices and values. Defaults to None. - is_coalesced (bool): Indicates whether the indices and values are coalesced. Defaults to False.

Methods: - is_coalesced(self): Check if the tensor is coalesced. - to(self, device: torch.DeviceObjType, non_blocking: bool = False): Move the tensor to the specified device. - diag(self, dims: Optional[Iterable[int]], return_sparse: bool = False): Extract diagonal elements from the tensor. The dimensions in dims will be take diagonal and put at dims[0] - sum(self, dims: Union[int, Optional[Iterable[int]]], return_sparse: bool = False): Compute the sum of tensor values along specified dimensions. return_sparse=True will return a sparse tensor, otherwise return a dense tensor. - max(self, dims: Union[int, Optional[Iterable[int]]], return_sparse: bool = False): Compute the maximum of tensor values along specified dimensions. return_sparse=True will return a sparse tensor, otherwise return a dense tensor. - mean(self, dims: Union[int, Optional[Iterable[int]]], return_sparse: bool = False): Compute the mean of tensor values along specified dimensions. return_sparse=True will return a sparse tensor, otherwise return a dense tensor. - unpooling(self, dims: Union[int, Iterable[int]], tarX): Perform unpooling operation along specified dimensions. - tuplewiseapply(self, func: Callable[[Tensor], Tensor]): Apply a function to each element of the tensor. - diagonalapply(self, func: Callable[[Tensor, LongTensor], Tensor]): Apply a function to diagonal elements of the tensor. - add(self, tarX, samesparse: bool): Add two sparse tensors together. samesparse=True means that two sparse tensor have the indice and can add values directly. - catvalue(self, tarX, samesparse: bool): Concatenate values of two sparse tensors. samesparse=True means that two sparse tensor have the indice and can cat values along the first dimension directly. - from_torch_sparse_coo(cls, A: torch.Tensor): Create a SparseTensor from a torch sparse COO tensor. - to_torch_sparse_coo(self) -> Tensor: Convert the SparseTensor to a torch sparse COO tensor.

Attributes: - indices (LongTensor): The indices of the sparse tensor. - values (Tensor): The values associated with the indices. - sparse_dim (int): The number of dimensions represented by the indices. - nnz (int): The number of non-zero values. - shape (torch.Size): The shape of the tensor. - sparseshape (torch.Size): The shape of the tensor up to the sparse dimensions. - denseshape (torch.Size): The shape of the tensor after the sparse dimensions.

add(tarX, samesparse: bool)[source]
catvalue(tarX, samesparse: bool)[source]
property denseshape
diag(dims: Iterable[int] | None, return_sparse: bool = False)[source]

TODO: unit test ??

diagonalapply(func: Callable[[Tensor, LongTensor], Tensor])[source]
classmethod from_torch_sparse_coo(A: Tensor)[source]
property indices
is_coalesced()[source]
max(dims: int | Iterable[int] | None, return_sparse: bool = False)[source]
mean(dims: int | Iterable[int] | None, return_sparse: bool = False)[source]
property nnz
property shape
property sparse_dim
property sparseshape
sum(dims: int | Iterable[int] | None, return_sparse: bool = False)[source]
to(device: DeviceObjType, non_blocking: bool = False)[source]
to_torch_sparse_coo() → Tensor[source]
tuplewiseapply(func: Callable[[Tensor], Tensor])[source]
unpooling(dims: int | Iterable[int], tarX)[source]

unpooling to of tarX indice dims: of tarX

unpooling_fromdense1dim(dims: int, X: Tensor)[source]

unpooling to of self shape. Note the dims is for self to maintain, and expand other dims

property values
pygho.backend.SpTensor.coalesce(edge_index: LongTensor, edge_attr: Tensor | None = None, reduce: str = 'sum') → Tuple[Tensor, Tensor | None][source]

Coalesces and reduces duplicate entries in edge indices and attributes.

Args:

  • edge_index (LongTensor): The edge indices.

  • edge_attr (Tensor or List[Tensor], optional): Edge weights or multi-dimensional edge features. If given as a list, it will be reshuffled and duplicates will be removed for all entries. (default: None)

  • reduce (str, optional): The reduction operation to use for merging edge features. Options include ‘sum’, ‘mean’, ‘min’, ‘max’, ‘mul’. (default: ‘sum’)

Returns:

  • Tuple[Tensor, Optional[Tensor]]: A tuple containing the coalesced edge indices and the coalesced and reduced edge attributes (if provided). If edge_attr is None, the second element will be None.

pygho.backend.SpTensor.decodehash(indhash: LongTensor, sparse_dim: int) → LongTensor[source]

Decodes a hashed LongTensor into tuples of indices.

This function takes a hashed LongTensor and decodes it into pairs of indices, which is commonly used in sparse tensor operations.

Parameters:

  • indhash (LongTensor): The input hashed LongTensor of shape (nnz).

  • sparse_dim (int): The number of dimensions represented by the hash.

Returns:

  • LongTensor: A LongTensor representing pairs of indices.

Raises:

  • AssertionError: If the input tensor doesn’t have the expected shape or if the sparse dimension is invalid.

Example:

indices = torch.tensor([[1, 2, 3], [4, 5, 6]], dtype=torch.long)
hashed = indicehash(indices)
indices = decodehash(hashed)
pygho.backend.SpTensor.decodehash_tight(indhash: LongTensor, dimsize: LongTensor) → LongTensor[source]

Decodes a tightly hashed LongTensor into pairs of indices considering dimension sizes.

Parameters: - indhash (LongTensor): The input hashed LongTensor of shape (nnz). - dimsize (LongTensor): The sizes of each dimension in the sparse tensor of shape (sparse_dim).

Returns: - LongTensor: A LongTensor representing pairs of indices.

Raises: - AssertionError: If the input tensors don’t have the expected shapes or if the total size exceeds the range that torch.long can express.

Example:

indices = torch.tensor([[1, 2, 0], [4, 1, 2]], dtype=torch.long)
dim_sizes = torch.tensor([3, 5], dtype=torch.long)
hashed = indicehash_tight(indices, dim_sizes)
indices = decodehash_tight(hashed, dim_sizes)
pygho.backend.SpTensor.indicehash(indice: LongTensor) → LongTensor[source]

Hashes a indice of shape (sparse_dim, nnz) to a single LongTensor of shape (nnz). Keep lexicographic order.

Parameters: - indice (LongTensor): The input indices tensor of shape (sparse_dim, nnz).

Returns: - LongTensor: A single LongTensor representing the hashed values.

Raises: - AssertionError: If the input tensor doesn’t have the expected shape or if the indices are too large or if there exists negative indice.

Example:

indices = torch.tensor([[1, 2, 3], [4, 5, 6]], dtype=torch.long)
hashed = indicehash(indices)
pygho.backend.SpTensor.indicehash_tight(indice: LongTensor, dimsize: LongTensor) → LongTensor[source]

Hashes a 2D LongTensor of indices tightly into a single LongTensor. Equivalently, it compute the indice of flattened sparse tensor with indice and dimsize

Parameters: - indice (LongTensor): The input indices tensor of shape (sparse_dim, nnz). - dimsize (LongTensor): The sizes of each dimension in the sparse tensor of shape (sparse_dim).

Returns: - LongTensor: A single LongTensor representing the tightly hashed values.

Raises: - AssertionError: If the input tensors don’t have the expected shapes or if the indices exceed the dimension sizes.

Example:

indices = torch.tensor([[1, 2, 0], [4, 1, 2]], dtype=torch.long)
dim_sizes = torch.tensor([3, 5], dtype=torch.long)
hashed = indicehash_tight(indices, dim_sizes)

pygho.backend.Spmamm module

pygho.backend.Spmamm.spmamm(A: SparseTensor, dim1: int, B: MaskedTensor, dim2: int, mask: BoolTensor | None = None, aggr: str = 'sum') → MaskedTensor[source]

SparseTensor-MaskedTensor multiplication.

This function performs multiplication between a SparseTensor A and a MaskedTensor B. The specified dimensions dim1 and dim2 are contracted during the multiplication, and the result is returned as a MaskedTensor.

Args:

  • A (SparseTensor): The SparseTensor with shape (B, n, m, *shape).

  • dim1 (int): The dimension to contract in the SparseTensor A.

  • B (MaskedTensor): The MaskedTensor with shape (B, m, *shape).

  • dim2 (int): The dimension to contract in the MaskedTensor B.

  • mask (BoolTensor, optional): The mask to apply to the resulting MaskedTensor. Default is None.

  • aggr (str, optional): The aggregation method for reduction during multiplication (e.g., “sum”, “max”). Default is “sum”.

Returns:

  • MaskedTensor: A new MaskedTensor with shape (B, n,*denseshapeshape) and the specified mask.

Notes: - This function performs multiplication between a SparseTensor and a MaskedTensor, contracting the specified dimensions. - The aggr parameter controls the reduction operation during multiplication. - The result is returned as a MaskedTensor.

pygho.backend.Spmm module

pygho.backend.Spmm.spmm(A: SparseTensor, dim1: int, X: Tensor, aggr: str = 'sum') → Tensor[source]

SparseTensor, Tensor matrix multiplication.

This function performs a matrix multiplication between a SparseTensor A and a dense tensor X along the specified dimension dim1. The result is a dense tensor. The aggr parameter specifies the reduction operation used for merging the resulting values.

Args:

  • A (SparseTensor): The SparseTensor used for multiplication.

  • dim1 (int): The dimension along which A is reduced.

  • X (Tensor): The dense tensor to be multiplied with A. It dim 0 will be reduced.

  • aggr (str, optional): The reduction operation to use for merging edge features (“sum”, “min”, “max”, “mean”). Defaults to “sum”.

Returns:

  • Tensor: A dense tensor containing the result of the matrix multiplication.

Notes:

  • A should be a 2-dimensional SparseTensor.

  • The dense shapes of A and X other than dim1 must be broadcastable.

pygho.backend.Spspmm module

pygho.backend.Spspmm.filterind(tar_ind: LongTensor, ind: LongTensor, bcd: LongTensor) → LongTensor[source]

A combination of Hadamard and Sparse Matrix Multiplication.

Given the indices tar_ind of sparse tensor A, the indices ind of sparse tensor BC, and the index array bcd, this function returns an index array acd, where (A ⊙ (BC)).val[a] = A.val[a] * scatter(B.val[c] * C.val[d], a).

Args:

  • tar_ind (LongTensor): The indices of sparse tensor A.

  • ind (LongTensor): The indices of sparse tensor BC.

  • bcd (LongTensor): An index array representing (BC).val.

Returns:

  • LongTensor: An index array acd representing the filtered indices.

Example:

tar_ind = torch.tensor([[0, 1, 1, 2],
                        [2, 1, 0, 2]], dtype=torch.long)
ind = torch.tensor([[2, 1, 0, 1],
                    [1, 0, 2, 2]], dtype=torch.long)
bcd = torch.tensor([[3, 2, 1, 0],
                    [6, 5, 4, 3],
                    [9, 8, 7, 6]], dtype=torch.long)
acd = filterind(tar_ind, ind, bcd)
pygho.backend.Spspmm.ptr2batch(ptr: LongTensor, dim_size: int) → LongTensor[source]

Converts a pointer tensor to a batch tensor. TODO: use torch_scatter gather instead?

This function takes a pointer tensor ptr and a dim_size and converts it to a batch tensor where each element in the batch tensor corresponds to a range of indices in the original tensor.

Args:

  • ptr (LongTensor): The pointer tensor, where ptr[0] = 0 and torch.all(diff(ptr) >= 0) is true.

  • dim_size (int): The size of the target dimension.

Returns:

  • LongTensor: A batch tensor of shape (dim_size,) where batch[ptr[i]:ptr[i+1]] = i.

pygho.backend.Spspmm.spsphadamard(A: SparseTensor, B: SparseTensor, b2a: LongTensor | None = None) → SparseTensor[source]

Element-wise Hadamard product between two SparseTensors.

This function performs the element-wise Hadamard product between two SparseTensors, A and B. The b2a parameter is an optional auxiliary index produced by the spsphadamard_ind function.

Args:

  • A (SparseTensor): The first SparseTensor.

  • B (SparseTensor): The second SparseTensor.

  • b2a (LongTensor, optional): An optional index array produced by spsphadamard_ind. If not provided, it will be computed.

Returns:

  • SparseTensor: A SparseTensor containing the result of the Hadamard product.

Notes:

  • Both A and B must be coalesced SparseTensors.

  • The dense shapes of A and B must be broadcastable.

pygho.backend.Spspmm.spsphadamard_ind(tar_ind: LongTensor, ind: LongTensor) → LongTensor[source]

Auxiliary function for SparseTensor-SparseTensor Hadamard product.

This function is an auxiliary function used in the Hadamard product of two sparse tensors. Given the indices tar_ind of sparse tensor A and the indices ind of sparse tensor B, this function returns an index array b2a of shape (ind.shape[1],) such that ind[:, i] matches tar_ind[:, b2a[i]] for each i. If b2a[i] is less than 0, it means ind[:, i] is not matched.

Args:

  • tar_ind (LongTensor): The indices of sparse tensor A.

  • ind (LongTensor): The indices of sparse tensor B.

Returns:

  • LongTensor: An index array b2a representing the matching indices between tar_ind and ind. b2a of shape ind.shape[1]. ind[:, i] matches tar_ind[:, b2a[i]]. if b2a[i]<0, ind[:, i] is not matched

Example:

tar_ind = torch.tensor([[0, 1, 1, 2],
                        [2, 1, 0, 2]], dtype=torch.long)
ind = torch.tensor([[2, 1, 0, 1],
                    [1, 0, 2, 2]], dtype=torch.long)
b2a = spsphadamard_ind(tar_ind, ind)
pygho.backend.Spspmm.spspmm(A: SparseTensor, dim1: int, B: SparseTensor, dim2: int, aggr: str = 'sum', bcd: LongTensor | None = None, tar_ind: LongTensor | None = None, acd: LongTensor | None = None) → SparseTensor[source]

SparseTensor SparseTensor matrix multiplication at a specified sparse dimension.

This function performs matrix multiplication between two SparseTensors, A and B, at the specified sparse dimensions dim1 and dim2. The result is a SparseTensor containing the result of the multiplication. The aggr parameter specifies the reduction operation used for merging the resulting values.

Args:

  • A (SparseTensor): The first SparseTensor.

  • dim1 (int): The dimension along which A is multiplied.

  • B (SparseTensor): The second SparseTensor.

  • dim2 (int): The dimension along which B is multiplied.

  • aggr (str, optional): The reduction operation to use for merging edge features (“sum”, “min”, “max”, “mean”). Defaults to “sum”.

  • bcd (LongTensor, optional): An optional auxiliary index array produced by spspmm_ind.

  • tar_ind (LongTensor, optional): An optional target index array for the output. If not provided, it will be computed.

  • acd (LongTensor, optional): An optional auxiliary index array produced by filterind.

Returns:

  • SparseTensor: A SparseTensor containing the result of the matrix multiplication.

Notes:

  • Both A and B must be coalesced SparseTensors.

  • The dense shapes of A and B must be broadcastable.

  • This function allows for optional indices bcd and tar_ind for improved performance and control.

pygho.backend.Spspmm.spspmm_ind(ind1: LongTensor, dim1: int, ind2: LongTensor, dim2: int, is_k2_sorted: bool = False) → Tuple[LongTensor, LongTensor][source]

Sparse-sparse matrix multiplication for indices.

This function performs a sparse-sparse matrix multiplication for indices. Given two sets of indices ind1 and ind2, this function eliminates dim1 in ind1 and dim2 in ind2, and concatenates the remaining dimensions.

The result represents the product of the input indices.

Args:

  • ind1 (LongTensor): The indices of the first sparse tensor of shape (sparsedim1, M1).

  • dim1 (int): The dimension to eliminate in ind1.

  • ind2 (LongTensor): The indices of the second sparse tensor of shape (sparsedim2, M2).

  • dim2 (int): The dimension to eliminate in ind2.

  • is_k2_sorted (bool, optional): Whether ind2 is sorted along dim2. Defaults to False.

Returns:

  • tarind: LongTensor: The resulting indices after performing the sparse-sparse matrix multiplication.

  • bcd: LongTensor: In tensor perspective (*i_1, k, *i_2), (*j_1, k, *j_2) -> (*i_1, *i_2, *j_1, *j_2). The return indice is of shape (3, nnz), (b, c, d), c represent index of *i, d represent index of *j, b represent index of output.For i=1,2,…,nnz, val1[c[i]] * val2[d[i]] will be add to output val’s b[i]-th element.

Example:

ind1 = torch.tensor([[0, 1, 1, 2],
                    [2, 1, 0, 2]], dtype=torch.long)
dim1 = 0
ind2 = torch.tensor([[2, 1, 0, 1],
                    [1, 0, 2, 2]], dtype=torch.long)
dim2 = 1
result = spspmm_ind(ind1, dim1, ind2, dim2)
pygho.backend.Spspmm.spspmpnn(A: SparseTensor, dim1: int, B: SparseTensor, dim2: int, C: SparseTensor, acd: LongTensor, message_func: Callable[[Tensor, Tensor, Tensor, LongTensor], Tensor], aggr: str = 'sum') → SparseTensor[source]

SparseTensor SparseTensor matrix multiplication at a specified sparse dimension using a message function.

This function extend matrix multiplication between two SparseTensors, A and B, at the specified sparse dimensions dim1 and dim2, while using a message function message_func to compute the messages sent from A to B and C. The result is a SparseTensor containing the result of the multiplication. The aggr parameter specifies the reduction operation used for merging the resulting values.

Args:

  • A (SparseTensor): The first SparseTensor.

  • dim1 (int): The dimension along which A is multiplied.

  • B (SparseTensor): The second SparseTensor.

  • dim2 (int): The dimension along which B is multiplied.

  • C (SparseTensor): The third SparseTensor, providing the target indice

  • acd (LongTensor): The auxiliary index array produced by a previous operation.

  • message_func (Callable): A callable function that computes the messages between A, B, and C.

  • aggr (str, optional): The reduction operation to use for merging edge features (“sum”, “min”, “max”, “mul”, “any”). Defaults to “sum”.

Returns:

  • SparseTensor: A SparseTensor containing the result of the matrix multiplication.

Notes:

  • Both A and B must be coalesced SparseTensors.

  • The dense shapes of A, B, and C must be broadcastable.

  • The message_func should take four arguments: A_values, B_values, C_values, and acd, and return messages based on custom logic.

pygho.backend.utils module

pygho.backend.utils.torch_scatter_reduce(dim: int, src: Tensor, ind: LongTensor, dim_size: int, aggr: str) → Tensor[source]

Applies a reduction operation to scatter elements from src to dim_size locations based on the indices in ind.

This function is a wrapper for torch.Tensor.scatter_reduce_ and is designed to scatter elements from src to dim_size locations based on the specified dimension dim and the indices in ind. The reduction operation is specified by the aggr parameter, which can be ‘sum’, ‘mean’, ‘min’, ‘max’.

Args:

  • dim (int): The dimension along which to scatter elements (only dim=0 is currently supported).

  • src (Tensor): The source tensor of shape (nnz, denseshape).

  • ind (LongTensor): The indices tensor of shape (nnz).

  • dim_size (int): The size of the target dimension for scattering.

  • aggr (str): The reduction operation to apply (‘sum’, ‘mean’, ‘min’, ‘max’, ‘mul’, ‘any’).

Returns:

  • Tensor: A tensor of shape (dim_size, denseshape) resulting from the scatter operation.

Raises:

  • AssertionError: If dim is not 0, or if ind is not 1-dimensional.

Example:

::

src = torch.tensor([[1, 2], [4, 5], [7, 8], [9, 10]], dtype=torch.float) ind = torch.tensor([2, 2, 0, 1], dtype=torch.long) dim_size = 3 aggr = ‘sum’ result = torch_scatter_reduce(0, src, ind, dim_size, aggr)

Module contents