tf.sparse.sampled_addmm

Performs the sampled matrix multiplication of two dense matrices.

tf.sparse.sampled_addmm(
    indices,
    values,
    dense_shape,
    mat1,
    mat2,
    beta=1.0,
    alpha=1.0,
    output_type=tf.dtypes.float32
)

Multiplies matrix mat1 by matrix mat2 at the locations defined by indices. The product is scaled and added to values, producing alpha * (mat1 @ mat2) * spy(indices) + beta * values.

The function spy(indices) is the sparsity pattern matrix derived from indices.

The mat1 and mat2 inputs must be tensors of rank >= 2 where the inner 2 dimensions specify valid matrix multiplication dimensions, and any further dimensions specify matching batch size.

The indices, values, and dense_shape inputs make up the components of a SparseTensor which defines the sparsity pattern of the output. The sparsity pattern has values of 1 at the positions defined by the SparseTensor, and 0 elsewhere.

The alpha and beta inputs are the scaling factors.

The supported types for values, mat1, and mat2 are: bfloat16, float16, float32, float64.

A simple 2-D tensor operation:

indices = tf.constant([0, 0, 1, 1], shape=[2, 2])
indices
<tf.Tensor: shape=(2, 2), dtype=int32, numpy=
array([[0, 0],
       [1, 1]], dtype=int32)>
values = tf.constant([0.5, 0.3])
values
<tf.Tensor: shape=(2,), dtype=float32, numpy=array([0.5, 0.3], dtype=float32)>
dense_shape = tf.constant([2, 2])
dense_shape
<tf.Tensor: shape=(2,), dtype=int32, numpy=array([2, 2], dtype=int32)>
mat1 = tf.constant([1, 2, 3, 4, 5, 6], shape=[2, 3], dtype=tf.float32)
mat1
<tf.Tensor: shape=(2, 3), dtype=float32, numpy=
array([[1., 2., 3.],
       [4., 5., 6.]], dtype=float32)>
mat2 = tf.constant([7, 8, 9, 10, 11, 12], shape=[3, 2], dtype=tf.float32)
mat2
<tf.Tensor: shape=(3, 2), dtype=float32, numpy=
array([[ 7.,  8.],
       [ 9., 10.],
       [11., 12.]], dtype=float32)>
tf.sparse.sampled_addmm(indices, values, dense_shape, mat1, mat2,
alpha=0.75, beta=0.25)
(<tf.Tensor: shape=(2, 2), dtype=int32, numpy=
array([[0, 0],
       [1, 1]], dtype=int32)>, <tf.Tensor: shape=(2,), dtype=float32, numpy=
array([ 43.625, 115.575], dtype=float32)>,
<tf.Tensor: shape=(2,), dtype=int32, numpy=array([2, 2], dtype=int32)>)

A batch operation:

indices = tf.constant([0, 1, 1, 0, 0, 0, 1, 0], shape=[2, 2, 2])
indices
<tf.Tensor: shape=(2, 2, 2), dtype=int32, numpy=
array([[[0, 1],
        [1, 0]],
       [[0, 0],
        [1, 0]]], dtype=int32)>
values = tf.constant([3, 5, 2, 7], shape=[2, 2], dtype=tf.float32)
values
<tf.Tensor: shape=(2, 2), dtype=float32, numpy=
array([[3., 5.],
       [2., 7.]], dtype=float32)>
dense_shape = tf.constant([2, 2])
dense_shape
<tf.Tensor: shape=(2,), dtype=int32, numpy=array([2, 2], dtype=int32)>
mat1 = tf.constant(np.arange(1, 13), shape=[2, 2, 3], dtype=tf.float32)
mat1
<tf.Tensor: shape=(2, 2, 3), dtype=float32, numpy=
array([[[ 1.,  2.,  3.],
        [ 4.,  5.,  6.]],
       [[ 7.,  8.,  9.],
        [10., 11., 12.]]], dtype=float32)>
mat2 = tf.constant(np.arange(13, 25), shape=[2, 3, 2], dtype=tf.float32)
mat2
<tf.Tensor: shape=(2, 3, 2), dtype=float32, numpy=
array([[[13., 14.],
        [15., 16.],
        [17., 18.]],
       [[19., 20.],
        [21., 22.],
        [23., 24.]]], dtype=float32)>
tf.sparse.sampled_addmm(indices, values, dense_shape, mat1, mat2,
alpha=0.75, beta=0.25)
(<tf.Tensor: shape=(2, 2, 2), dtype=int32, numpy=
array([[[0, 1],
        [1, 0]],
       [[0, 0],
        [1, 0]]], dtype=int32)>, <tf.Tensor: shape=(2, 2), dtype=float32,
numpy=array([[ 75.75, 173.  ],
       [381.5 , 524.5 ]], dtype=float32)>,
<tf.Tensor: shape=(2,), dtype=int32, numpy=array([2, 2], dtype=int32)>)

indices tf.Tensor containing coordinates for the rows and columns to be multiplied. Must have rank > 1.
values tf.Tensor containing the values to be scaled and added to the sampled dot product.
dense_shape tf.Tensor defining the dense shape of the output.
mat1 tf.Tensor to be multiplied. Must have rank > 1.
mat2 tf.Tensor to be multiplied. Must have rank > 1.
beta Number to be multiplied with values. Defaults to 1.0.
alpha Number to be multiplied with the sampled dot product of mat1 and mat2. Defaults to 1.0.
output_type The output datatype if needed. Defaults to float32.

A tuple representing the SparseTensor components of the result of the operation.

ValueError If dense_shape does not match the shape of the product.