Permutation of Sparse Tensor

[ignaciogavier]

🚀 The feature, motivation and pitch

I'm working with COO sparse tensors and would like to get a permutation of any COO sparse tensor.
In the current version, torch.permute throws the following error:
RuntimeError: sparse tensors do not have strides
Also, in-place modification of indices is not possible.

Alternatives

I have considered the following solution (it wouldn't be the same as torch.permute since this one returns a view of the original tensor):

torch.permute_sparse(input, dims) → Tensor
Returns a copy of the original tensor input with its dimensions permuted.

• input (Tensor) – the input tensor.
• dims (tuple of python:ints) – The desired ordering of dimensions

The source code:

def permute_sparse(input, dims):
    dims = torch.LongTensor(dims)
    return torch.sparse_coo_tensor(indices=input._indices()[dims], values=input._values(), size=torch.Size(torch.tensor(input.size())[dims]))

An example:

>>> indices = torch.tensor([[1, 1],
                            [2, 1],
                            [1, 0],
                            [3, 3]])
>>> values = torch.tensor([ 1, -1])
>>> s = torch.sparse_coo_tensor(indices=indices, values=values, size=(2, 3, 4, 5))
>>> s.size()
torch.Size([2, 3, 4, 5])
>>> permute_sparse(s, (3, 2, 0, 1)).size()
torch.Size([5, 4, 2, 3])

Additional context

No response

cc @nikitaved @pearu @cpuhrsch @amjames

[cpuhrsch]

Hello @ignaciogavier,

Thank you for your feature request. If you can share it, could I ask you to add some more detail on your use case and the model you're trying to use this for?

Thanks,
Christian

[ignaciogavier]

Hi @cpuhrsch,
Thanks for your response.

I'm working with a representation of (sparse) multilinear polynomials and I'm using sparse tensors. I use the indices order to reference the order of the variables in the polynomial and the values to store the coefficient of the polynomial terms.

As an example, addition is a common operation between two polynomials, but before performing it, the order of the variables have to be "aligned". That is, the indices should correspond to the same variable.
Example (p3 = p1 + p2):

>>> p1(a,b,c) = a + 2c - 3ab   ## (the tensor values are [1 2 -3], while the indices are [[1 0 1], [0 0 1], [0 1 0]]) 
>>> p2(b,c,a) = a + c          ## (the tensor values are [1 1], while the indices are [[0 0], [0 1], [1 0]]; note that c is the second variable so its index list is the second one in indices) 
>>> p3(a,b,c) = 2a + 3c - 3ab  ## (the tensor values are [2 3 -3], while the indices are [[1 0 1], [0 0 1], [0 1 0]])

Note that the tensor of p3 can only be computed as p1 + p2 if we first permute p2 according to the permutation (b,c,a)->(a,b,c).

In addition, any operation between the polynomials requires that the order of indices match for each one of them.

Best,
Ignacio

[cpuhrsch]

Thanks for the details @ignaciogavier,

Could you go into a bit more detail how the torch.Tensors you're using here are constructed? For example, are you constructing 2D torch.Tensors that represent a Polynomial matrix? Such a matrix could then be compressed using a sparse layout such as COO or CSR.

Thanks,
Christian

[ignaciogavier]

Thanks for the link to Polynomial matrix! It is similar but with some differences. I'm using multilinear polynomials (all variables are powered to 0 or 1), but that's a particular case of multivariable polynomials.

To construct the matrix representation of the multilinear polynomial, take the following example:

p(a,b,c) = a + 2c - 3ab = 0 + a + 0b + 2c + (-3)ab + 0bc + 0ac + 0abc = 0 + a + 0b + (-3)ab + 2c + 0ac + 0bc + 0abc

In a general way, a polynomial that depends on n variables, has 2^n terms. So the coefficients can be arranged in a multidimensional tensor of size (2,...,2), i.e. n dimensions with 2 coordinates each --here it's 2 because the variable is powered to either 0 or 1 but this is generalizable to non-multilinear: you'd have to have d+1 coordinates in each dimension, being d the degree in that dimension--.
In the previous example, that would be

[[[0 1],      # coefficients for constant term and term "a"
  [0 -3]],      # coefficients for term "b" and term "ab"
 [[2 0],      # coefficients for term "c" and term "ac"
  [0 0]]]      # coefficients for term "bc" and term "abc"

Now the sparse tensor idea came because I'm working with very long polynomials where most of the coefficients are zero.

[cpuhrsch]

@ignaciogavier - Ah I see! Ok, thanks for the additional detail.