TorchUST#

class nvmath.sparse.ust.interfaces.torch_interface.TorchUST( shape: Size, device: device, dtype: dtype, layout: layout, requires_grad: bool, ust: Tensor, transposed: bool, )[source]#

This class wraps the universal sparse tensor as a torch.Tensor.

The purpose of this class is to inject the UST into PyTorch models with minimal code changes. Objects of the TorchUST class behave like PyTorch tensors, but transparently use an UST implementation for the actual operations. A set of operations commonly used in models is supported (otherwise a “can’t perform” message is prompted; please let us know of such cases, so we can extend the supported set accordingly).

Examples

>>> import torch
>>> from nvmath.sparse.ust.interfaces.torch_interface import TorchUST

Create two torch vectors.

>>> x = (1.0 + torch.arange(32)).cuda()
>>> y = (2.0 + torch.arange(32)).cuda()

Now perform a dot product operation.

>>> z = torch.dot(x, y)

Convert the first operand to UST.

>>> x = TorchUST.from_torch(x)

Now perform the same dot production operation. It transparently uses the UST implementation!

>>> z = torch.dot(x, y)

classmethod from_torch( original_tensor: Tensor, ) → TorchUST[source]#: Constructs a TorchUST from the given torch.Tensor.