torch.sparse_bsr_tensor#

torch.sparse_bsr_tensor(crow_indices,col_indices,values,size=None,*,dtype=None,device=None,pin_memory=False,requires_grad=False,check_invariants=None)→Tensor#

Constructs asparse tensor in BSR (Block Compressed Sparse Row)) with specified 2-dimensional blocks at the givencrow_indices andcol_indices. Sparse matrixmultiplication operations in BSR format are typically faster than thatfor sparse tensors in COO format. Make you have a look atthenote on the data type of the indices.

Note

If thedevice argument is not specified the device of the givenvalues and indices tensor(s) must match. If, however, theargument is specified the input Tensors will be converted to thegiven device and in turn determine the device of the constructedsparse tensor.

Parameters

• crow_indices (array_like) – (B+1)-dimensional array of size(*batchsize,nrowblocks+1). The last element of eachbatch is the number of non-zeros. This tensor encodes theblock index in values and col_indices depending on where thegiven row block starts. Each successive number in the tensorsubtracted by the number before it denotes the number ofblocks in a given row.

• col_indices (array_like) – Column block coordinates of each blockin values. (B+1)-dimensional tensor with the same length asvalues.

• values (array_list) – Initial values for the tensor. Can be a list,tuple, NumPyndarray, scalar, and other types thatrepresents a (1 + 2 + K)-dimensional tensor whereK is thenumber of dense dimensions.

• size (list, tuple,torch.Size, optional) – Size of thesparse tensor:(*batchsize,nrows*blocksize[0],ncols*blocksize[1],*densesize) whereblocksize==values.shape[1:3]. If not provided, the size will beinferred as the minimum size big enough to hold all non-zeroblocks.

Keyword Arguments

• dtype (torch.dtype, optional) – the desired data type ofreturned tensor. Default: if None, infers data type fromvalues.

• device (torch.device, optional) – the desired device ofreturned tensor. Default: if None, uses the current devicefor the default tensor type (seetorch.set_default_device()).device will bethe CPU for CPU tensor types and the current CUDA device forCUDA tensor types.

• pin_memory (bool,optional) – If set, returned tensor would be allocated inthe pinned memory. Works only for CPU tensors. Default:False.

• requires_grad (bool,optional) – If autograd should record operations on thereturned tensor. Default:False.

• check_invariants (bool,optional) – If sparse tensor invariants are checked.Default: as returned bytorch.sparse.check_sparse_tensor_invariants.is_enabled(),initially False.

Example:

>>>crow_indices=[0,1,2]>>>col_indices=[0,1]>>>values=[[[1,2],[3,4]],[[5,6],[7,8]]]>>>torch.sparse_bsr_tensor(torch.tensor(crow_indices,dtype=torch.int64),...torch.tensor(col_indices,dtype=torch.int64),...torch.tensor(values),dtype=torch.double)tensor(crow_indices=tensor([0, 1, 2]),       col_indices=tensor([0, 1]),       values=tensor([[[1., 2.],                       [3., 4.]],                      [[5., 6.],                       [7., 8.]]]), size=(2, 2), nnz=2, dtype=torch.float64,       layout=torch.sparse_bsr)