mindspore.RowTensor

class mindspore.RowTensor(indices, values, dense_shape)[source]

A sparse representation of a set of tensor slices at given indices.

An RowTensor is typically used to represent a subset of a larger tensor dense of shape [L0, D1, .. , DN] where L0 >> D0.

The values in indices are the indices in the first dimension of the slices that have been extracted from the larger tensor.

The dense tensor dense represented by an RowTensor slices has dense[slices.indices[i], :, :, :, …] = slices.values[i, :, :, :, …].

For example, if indices is [0], values is [[1, 2]], dense_shape is (3, 2), then the dense representation of the row tensor will be:

[[1, 2],
 [0, 0],
 [0, 0]]

RowTensor can only be used in the Cell’s construct method.

Note

RowTensor is not supported in pynative mode.

Parameters
  • indices (Tensor) – A 1-D integer Tensor of shape [D0].

  • values (Tensor) – A Tensor of any dtype of shape [D0, D1, …, Dn].

  • dense_shape (tuple(int)) – An integer tuple which contains the shape of the corresponding dense tensor.

Returns

RowTensor, composed of indices, values, and dense_shape.

Examples

>>> import mindspore as ms
>>> import mindspore.nn as nn
>>> from mindspore import Tensor, RowTensor
>>> class Net(nn.Cell):
...     def __init__(self, dense_shape):
...         super(Net, self).__init__()
...         self.dense_shape = dense_shape
...     def construct(self, indices, values):
...         x = RowTensor(indices, values, self.dense_shape)
...         return x.values, x.indices, x.dense_shape
>>>
>>> indices = Tensor([0])
>>> values = Tensor([[1, 2]], dtype=ms.float32)
>>> out = Net((3, 2))(indices, values)
>>> print(out[0])
[[1. 2.]]
>>> print(out[1])
[0]
>>> print(out[2])
(3, 2)