Tensor#

class nvmath.sparse.ust.Tensor(
extents,
*,
tensor_format=None,
index_type='int32',
dtype='float32',
package=None,
options=None,
stream=None,
)[source]#

The universal sparse tensor binds extents, data type, index type, etc with a universal sparse tensor format object.

Note

The constructor is currently private and for internal use only. Users should use any of the methods Tensor.from_package(), Tensor.from_file(), or Tensor.convert() to construct a UST.

Parameters:

  • extents – The extents (shape) of the sparse tensor as a Python sequence

  • tensor_format – The sparse tensor format as a TensorFormat object.

  • index_type – The index type of the tensor as a NumPy-recognizable string or CudaDataType.

  • dtype – The datatype of the tensor as a NumPy-recognizable string or CudaDataType.

  • package – The package using which the dense tensors underlying the sparse representation will be allocated. It must be a Python module object.

Attributes

base#

Get the “package” tensor from which the UST view was created, if available.

device#

The memory space of this UST object (“cpu” or “cuda”).

device_id#

The device ID of this UST object (CUDA ordinal or “cpu”).

dtype#

The data type of this UST object as a string.

extents#

The extents (shape) of this UST object.

index_type#

The index type of this UST object as a string.

levels#

The extents corresponding to each level for this UST object.

nse#

The number of stored elements in this UST object.

num_dimensions#

The number of dimensions of this UST object.

num_levels#

The number of levels of this UST object.

shape#

The extents (shape) of this UST object.

size#

The envelope of this UST object (product of the extents).

tensor_format#

The UST format as a nvmath.sparse.ust.TensorFormat object.

val#

The explicit values stored in this UST object.

Methods

animate(name=None)[source]#

Animates the tensor nonzero structure (for 3D). The method either saves the animated GIF in a file named name when name is not None, or otherwise returns a IPython.display.HTML that can be embedded by the caller.

This method scales to larger tensors.

Parameters:

name – filename to save the animated GIF to, if not None

Returns:

IPython.display.HTML, if name is None (to embed in other output)

Examples

>>> import torch
>>> from nvmath.sparse.ust import Tensor

Create a 3-D tensor, as COO.

>>> a = torch.eye(32).repeat(32, 1, 2).to_sparse()

Convert the tensor to UST.

>>> u = Tensor.from_package(a)

Obtain the animation as HTML.

>>> html = u.animate()
clone(stream=None)[source]#

Clone the UST. The sparse representation of the UST is copied to the target, while any associated kernel is shared between source and target.

Note

Any kernel associated with the source is shared with the target.

Parameters:

stream – Provide the CUDA stream to use for executing the operation. Acceptable inputs include cudaStream_t (as Python int), cupy.cuda.Stream, and torch.cuda.Stream. If a stream is not provided, the current stream from the operand package will be used. See Stream Semantics for more details on stream handling.

Returns:

A clone of this UST.

convert(*, tensor_format=None, index_type=None, dtype=None)[source]#

Convert a UST into a new UST with the specified format, index type, and data type. The default values of these for the target UST are taken from the corresponding ones in the source UST.

Please note that this is a proof-of-concept implementation that supports converting any UST format into any other UST format. Even though some fast-path solutions are provided, the current general solution in TensorConverter has not been optimized for speed yet.

Note

The target shares a source kernel only for trivial conversions.

Parameters:

  • tensor_format – The sparse tensor format as a TensorFormat object.

  • index_type – The index type of the tensor as a NumPy-recognizable string or CudaDataType.

  • dtype – The datatype of the tensor as a NumPy-recognizable string or CudaDataType.

Returns:

A UST in the specified format, with the specified index and data types.

copy_(src, stream=None)[source]#

Copy a compatible UST into this one, where compatible means that the source UST has the same shape, tensor format, non-zero structure, index type, and data type as this one.

Note

Any kernel associated with the source is not copied over to target, and any kernel that was associated with this one is reset to None.

Parameters:

  • src – The source UST (nvmath.sparse.ust.Tensor) that will be copied into this one.

  • stream – Provide the CUDA stream to use for executing the operation. Acceptable inputs include cudaStream_t (as Python int), cupy.cuda.Stream, and torch.cuda.Stream. If a stream is not provided, the current stream from the operand package will be used. See Stream Semantics for more details on stream handling.

crd(level)[source]#

The coordinates corresponding to the specified level.

Parameters:

level – An ordinal specifying the level.

draw(name=None)[source]#

Draws the tensor contents (for 1D, 2D, 3D). The method saves the constructed PIL.Image in a file named name when name is not None. It always returns the PIL.Image object directly as well, so that the caller can directly manipulate the constructed image in some other manner.

This method is useful to illustrate tensor contents of smaller examples.

Parameters:

name – filename to save the image to, if not None

Returns:

PIL.Image, can be displayed with show()

Examples

>>> import scipy.sparse as sp
>>> from nvmath.sparse.ust import Tensor

Create a sparse scipy matrix in CSR format.

>>> a = sp.random_array((8, 8), density=0.1, format="csr", dtype="float64")

Convert the scipy matrix to UST.

>>> u = Tensor.from_package(a)

Obtain the tensor contents as image (and e.g. display as img.show()).

>>> img = u.draw()
draw_raw(name=None)[source]#

Draws the tensor nonzero structure (for 2D, 3D). The method saves the constructed PIL.Image in a file named name when name is not None. It always returns the PIL.Image object directly as well, so that the caller can directly manipulate the constructed image in some other manner.

This method scales to larger tensors.

Parameters:

name – filename to save the image to, if not None

Returns:

PIL.Image, can be displayed with show()

Examples

>>> import scipy.sparse as sp
>>> from nvmath.sparse.ust import Tensor

Create a sparse scipy matrix in CSR format.

>>> a = sp.random_array((512, 512), density=0.01, format="csr", dtype="float64")

Convert the scipy matrix to UST.

>>> u = Tensor.from_package(a)

Obtain the tensor nonzero structure as image (e.g. display as img.show()).

>>> img = u.draw_raw()
draw_storage(name=None)[source]#

Draws the tensor storage (for any UST). The method saves the constructed PIL.Image in a file named name when name is not None. It always returns the PIL.Image object directly as well, so that the caller can directly manipulate the constructed image in some other manner.

This method is useful to illustrate the UST storage of smaller examples.

Parameters:

name – filename to save the image to, if not None

Returns:

PIL.Image, can be displayed with show()

Examples

>>> import scipy.sparse as sp
>>> from nvmath.sparse.ust import Tensor

Create a sparse scipy matrix in CSR format.

>>> a = sp.random_array((8, 8), density=0.1, format="csr", dtype="float64")

Convert the scipy matrix to UST.

>>> u = Tensor.from_package(a)

Obtain the tensor storage as image (and e.g. display as img.show()).

>>> img = u.draw_storage()
classmethod from_file(filename, stream=None)[source]#

A helper to create a universal sparse tensor (in COO representation) from a file. The supported file formats are the Matrix Market format (.mtx) for matrices and the FROSTT format (.tns) for tensors.

classmethod from_package(tensor, stream=None)[source]#

Create an universal sparse tensor from a sparse package tensor. This is a zero-copy operation where the UST shares a view of the sparse data with the sparse tensor. A strided tensor (like numpy.ndarray or torch.Tensor) can also be viewed as a UST, with the current limitation that there are no “holes” (the strided tensor is dense, in other words).

The currently supported libraries are NumPy, SciPy, CuPy, and PyTorch and both CPU and GPU (CUDA) memory spaces are supported.

Parameters:

  • tensor – The source sparse tensor from NumPy, SciPy, CuPy, or PyTorch packages.

  • stream – Provide the CUDA stream to use for executing the operation. Acceptable inputs include cudaStream_t (as Python int), cupy.cuda.Stream, and torch.cuda.Stream. If a stream is not provided, the current stream from the operand package will be used. See Stream Semantics for more details on stream handling.

Returns:

A UST view of the sparse or dense package tensor.

get_value(indices)[source]#

Retrieve the value at specified indices (coordinate) or None if there is no value at that coordinate. This is not random access, and searches compressed levels.

Returns:

tensor.get_value[i, j, k, ...]` returns the value of the element appearing at logical dimension indices ``i, j, k, ...` or ``None if there is no element at that coordinate.

pos(level)[source]#

The positions corresponding to the specified level.

Parameters:

level – An ordinal specifying the level.

run_kernel(stream=None)[source]#

Run the kernel associated with this UST object.

set_kernel(user_code, *, with_indices=True, arch=None)[source]#

Set the callback code (a unary Python function, possibly already compiled into LTO-IR) and link this with the actual traversal kernel for the UST. This can be used to apply an user-defined transformation to each element of the UST, where the user-defined transformation can potentially depend on the coordinates (indices) of the element.

Parameters:

  • code – The callback code as a unary Python function object or as a bytes object with the LTO-IR code. The function must take a single argument representing the original value if with_indices is False, otherwise it takes the original value followed by the coordinate for that value, as a sequence. The original value must be representable in the data type of the UST, while the indices must be integers.

  • with_indices – A flag to choose between the value-only or value-with-indices signatures.

Tip

If the indices (coordinate) are not needed for the user-defined transformation, it is much more efficient to set with_indices=False.

to(device_id: int | Literal['cpu'], stream=None)[source]#

Copy the UST to a different device (contents only).

The sparse representation is not copied (it is a view) if the UST is already on the requested device.

Note

Any kernel associated with the source is not copied over to target.

Parameters:

  • device_id – The CUDA device ordinal, or “cpu”.

  • stream – Provide the CUDA stream to use for executing the operation. Acceptable inputs include cudaStream_t (as Python int), cupy.cuda.Stream, and torch.cuda.Stream. If a stream is not provided, the current stream from the operand package will be used. See Stream Semantics for more details on stream handling.

Returns:

A copy of the UST on the specified device.

to_package(options=None, stream=None)[source]#

Create a sparse package tensor from the universal sparse tensor. This is a zero-copy operation where the UST shares a view of the sparse data with the package tensor.

Returns:

The sparse or dense package tensor view into this UST.