I am wondering about how to best implement certain types of structured matrices in Pytorch. By “structured matrix,” I mean matrices that can behave as a FloatTensor (for example) in every way, but certain operations can be performed more efficiently. For example, matrix-vector multiplications and (approximate) linear solves with Toeplitz matrices are very fast. Other examples might include triangular matrices (to represent LU or Cholesky factorizations), or matrices that are the Kronecker product of smaller matrices.
I would additionally like to use these structured matrices in autograd: ideally it would be possible for the backward computation of some function to return a structured matrix, so that when earlier nodes in the graph use this structured matrix in their computations, the operations are efficient.
My first thought about how to accomplish this was to create, for example, ToeplitzFloatTensor, that extends FloatTensor and rewrites a few of the methods to handle the added structure efficiently.
However, Pytorch does not seem to easily allow extending the base Tensor classes. Or, at least, instances of classes extending FloatTensor cannot be used in a Variable without receiving an error:
RuntimeError: Variable data has to be a tensor…
Furthermore, adding such a class to torch._tensor_classes simply results in an error later down the line.
Is extending the base Tensor classes at all possible in Pytorch, or is there some other preferred pattern for making Tensors that represent structured matrices?