sparse

construct_block_diag_coo_indices_and_shape(*block_shapes: tuple[int, int], device: ~torch.device = None) -> (<class 'torch.Tensor'>, tuple[int, int])

Construct the indices for a block diagonal sparse matrix in COOrdinate format.

>>> blocks = [torch.ones((2, 2)), torch.ones((3, 3))]
>>> indices, shape = construct_block_diag_coo_indices_and_shape(*[b.shape for b in blocks])
>>> torch.sparse_coo_tensor(indices, torch.cat([b.flatten() for b in blocks]), size=shape).to_dense()

Parameters:

• block_shapes – The shapes of the blocks.

• device – Which device to use.

Returns:

The indices of the blocks. tuple: The shape of the block diagonal matrix.

Return type:

torch.Tensor

construct_block_diag_coo_tensor(*blocks: Tensor) → Tensor

Construct a block diagonal tensor from the given blocks.

>>> blocks = [torch.ones((2, 2)), torch.ones((3, 3))]
>>> construct_block_diag_tensor(*blocks)

Parameters:

blocks – The blocks of the block diagonal tensor.

Returns:

The block diagonal tensor.

Return type:

torch.Tensor

construct_block_diag_coo_tensor_indices_and_shape_from_sparse(*blocks: ~torch.Tensor, device: ~torch.device = None) -> (<class 'torch.Tensor'>, tuple[int, int])

Construct the indices for a block diagonal sparse matrix in COOrdinate format where the blocks are sparse tensors.

Parameters:

• blocks – The blocks of the block diagonal tensor.

• device – Which device to use.

Returns:

The indices of the blocks. tuple: The shape of the block diagonal matrix.

Return type:

torch.Tensor