Working with sparse tensors

View on TensorFlow.org Run in Google Colab View on GitHub Download notebook

When working with tensors that contain a lot of zero values, it is important to store them in a space- and time-efficient manner. Sparse tensors enable efficient storage and processing of tensors that contain a lot of zero values. Sparse tensors are used extensively in encoding schemes like TF-IDF as part of data pre-processing in NLP applications and for pre-processing images with a lot of dark pixels in computer vision applications.

Sparse tensors in TensorFlow

TensorFlow represents sparse tensors through the tf.sparse.SparseTensor object. Currently, sparse tensors in TensorFlow are encoded using the coordinate list (COO) format. This encoding format is optimized for hyper-sparse matrices such as embeddings.

The COO encoding for sparse tensors is comprised of:

  • values: A 1D tensor with shape [N] containing all nonzero values.
  • indices: A 2D tensor with shape [N, rank], containing the indices of the nonzero values.
  • dense_shape: A 1D tensor with shape [rank], specifying the shape of the tensor.

A nonzero value in the context of a tf.sparse.SparseTensor is a value that's not explicitly encoded. It is possible to explicitly include zero values in the values of a COO sparse matrix, but these "explicit zeros" are generally not included when referring to nonzero values in a sparse tensor.

Creating a tf.sparse.SparseTensor

Construct sparse tensors by directly specifying their values, indices, and dense_shape.

import tensorflow as tf

2024-10-25 01:24:09.202320: E external/local_xla/xla/stream_executor/cuda/cuda_fft.cc:477] Unable to register cuFFT factory: Attempting to register factory for plugin cuFFT when one has already been registered
WARNING: All log messages before absl::InitializeLog() is called are written to STDERR
E0000 00:00:1729819449.223893   16549 cuda_dnn.cc:8310] Unable to register cuDNN factory: Attempting to register factory for plugin cuDNN when one has already been registered
E0000 00:00:1729819449.230517   16549 cuda_blas.cc:1418] Unable to register cuBLAS factory: Attempting to register factory for plugin cuBLAS when one has already been registered

st1 = tf.sparse.SparseTensor(indices=[[0, 3], [2, 4]],
                      values=[10, 20],
                      dense_shape=[3, 10])

W0000 00:00:1729819451.911465   16549 gpu_device.cc:2344] Cannot dlopen some GPU libraries. Please make sure the missing libraries mentioned above are installed properly if you would like to use GPU. Follow the guide at https://www.tensorflow.org/install/gpu for how to download and setup the required libraries for your platform.
Skipping registering GPU devices...

When you use the print() function to print a sparse tensor, it shows the contents of the three component tensors:

print(st1)

SparseTensor(indices=tf.Tensor(
[[0 3]
 [2 4]], shape=(2, 2), dtype=int64), values=tf.Tensor([10 20], shape=(2,), dtype=int32), dense_shape=tf.Tensor([ 3 10], shape=(2,), dtype=int64))

It is easier to understand the contents of a sparse tensor if the nonzero values are aligned with their corresponding indices. Define a helper function to pretty-print sparse tensors such that each nonzero value is shown on its own line.

def pprint_sparse_tensor(st):
  s = "<SparseTensor shape=%s \n values={" % (st.dense_shape.numpy().tolist(),)
  for (index, value) in zip(st.indices, st.values):
    s += f"\n  %s: %s" % (index.numpy().tolist(), value.numpy().tolist())
  return s + "}>"

print(pprint_sparse_tensor(st1))

<SparseTensor shape=[3, 10] 
 values={
  [0, 3]: 10
  [2, 4]: 20}>

You can also construct sparse tensors from dense tensors by using tf.sparse.from_dense, and convert them back to dense tensors by using tf.sparse.to_dense.

st2 = tf.sparse.from_dense([[1, 0, 0, 8], [0, 0, 0, 0], [0, 0, 3, 0]])
print(pprint_sparse_tensor(st2))

<SparseTensor shape=[3, 4] 
 values={
  [0, 0]: 1
  [0, 3]: 8
  [2, 2]: 3}>

st3 = tf.sparse.to_dense(st2)
print(st3)

tf.Tensor(
[[1 0 0 8]
 [0 0 0 0]
 [0 0 3 0]], shape=(3, 4), dtype=int32)

Manipulating sparse tensors

Use the utilities in the tf.sparse package to manipulate sparse tensors. Ops like tf.math.add that you can use for arithmetic manipulation of dense tensors do not work with sparse tensors.

Add sparse tensors of the same shape by using tf.sparse.add. 

st_a = tf.sparse.SparseTensor(indices=[[0, 2], [3, 4]],
                       values=[31, 2], 
                       dense_shape=[4, 10])

st_b = tf.sparse.SparseTensor(indices=[[0, 2], [3, 0]],
                       values=[56, 38],
                       dense_shape=[4, 10])

st_sum = tf.sparse