tf.sparse.sparse_dense_matmul can't broadcast as tf.linalg.matmul does.

[breadbread1984]

Please make sure that this is a feature request. As per our GitHub Policy, we only address code/doc bugs, performance issues, feature requests and build/installation issues on GitHub. tag:feature_template

System information

• TensorFlow version (you are using): 2.1.0
• Are you willing to contribute it (Yes/No): No

Describe the feature and the current behavior/state.

tf.sparse.sparse_dense_matmul can't broadcast on the leading dimensions. the following code is an example.

#!/usr/bin/python3

import numpy as np;
import tensorflow as tf;

def main():

    a = tf.sparse.SparseTensor(
        indices = [[0, 0, 1], [0, 0, 2], [0, 1, 2], [0, 1, 3], [0, 2, 1], [0, 2, 3]],
        values = [1., 1., 1., 1., 1., 1.],
        dense_shape = [1, 3, 4]
    ); # a.shape = (1,3,4)
    b = tf.constant(np.random.normal(size = (4, 4, 5)), dtype = tf.float32);
    c = tf.linalg.matmul(tf.sparse.to_dense(a),b); # will be succeed
    c = tf.sparse.sparse_dense_matmul(a,b); # will be fail

if __name__ == "__main__":

    main();

Will this change the current api? How?

no api need to be changed.

Who will benefit with this feature?

user handling sparse matrices will be benefit from the change.

Any Other info.

[familyguy12]

Hi, I would like to handle this and add the feature. Please let me know if no one is working on it.

[jvishnuvardhan]

@familyguy12 Thanks for coming forward to support the development. Please feel free to contribute through PR. Thanks!

[familyguy12]

@jvishnuvardhan It seems that this method is working as mentioned in the documentation. If I understand correctly, I am supposed to modify the way sparse matrix multiplication works so that it handles the case when the first arg is of rank != 2. Is that right?

[breadbread1984]

the tensorflow's traditional matmul does matrix production on the last two dimensions of the two inputs. if the leading dimensions of the two inputs are not equal, matmul may broadcast the small one to the large one if plausible.

[familyguy12]

In the code mentioned, tf.rank(a).eval() is 3. The documentation says that the sparse multiplication works only for rank-2 tensors. Even in the traditional matmul function, it says that

If one or both of the matrices contain a lot of zeros, a more efficient multiplication algorithm can be used by setting the corresponding a_is_sparse or b_is_sparse flag to True. These are False by default. This optimization is only available for plain matrices (rank-2 tensors) with datatypes bfloat16 or float32.

So, basically it is not supposed to work for matrices that are sparse, with rank > 2. IMO, it could work as the multiplication is performed on the last two dimensions only. So, I should handle the case where rank>2 in sparse_matmul and perform multiplication only on the last two dimensions, instead of giving a ValueError, right? This is my first contribution to the repository and I apologise for too many questions.

[breadbread1984]

the tf.sparse.sparse_dense_matmul can only handle dot product of two 2-rank tensors. the function is too constrained. I would appreciate if dot product between higher rank tensors are supported. besides, I would like to have auto broadcast of sparse_dense_matmul supported. Thx.

[breadbread1984]

if sparse_dense_matmul can't support dot product between tensors over rank 2. please add a tf.sparse.map_fn which support tensor slice over sparse tensor. Or some operations are impossible.

[breadbread1984]

please support matmul for sparse tensor. pytorch has deepspeed which can handle sparse tensor matmul having rank over 2. implementing sparse attention with tensorflow 2 is impossible without such feature.

[breadbread1984]

@familyguy12 any progress?

[sjtusmartboy]

@familyguy12 any progress?

[MPKenning]

I believe this is an important feature to enable. If I hadn't already programmed so much of my project in Tensorflow, I would have switched to PyTorch.

[chunduriv]

I was able to replicate the issue in tf-nightly 2.10.0-dev20220726. Please find the gist for reference. Thank you.

[dmitrybaltin]

I also vote for the feature. I was unpleasantly surprised when I realized that it's impossible to multiply sparse matrixes with a rank > 3 ( By the way, thanks for the info about pytorch. Didn't plan to switch to it, but might have to

[MPKenning]

Hello everyone,

There is no native satisfactory solution to performing a matrix multiplication between a SparseTensor and a Tensor and between two SparseTensors. So I made a solution. I think it runs well enough. Everyone is free to use it. I don't demand anything but acknowledgement in return.

I have posted the code on Colab.

Here's a link to a GitHub repository I made with the code (and a little more on the side).

EDIT: The function is not faster than a for-loop or dense matrix multiplication. I do not recommend it. The speed-up is not worth it. I suppose it works as a proof of concept. I welcome contributions.

[tilakrayal]

Hi,

Thank you for opening this issue. Since this issue has been open for a long time, the code/debug information for this issue may not be relevant with the current state of the code base.

The Tensorflow team is constantly improving the framework by fixing bugs and adding new features. We suggest you try the latest TensorFlow version with the latest compatible hardware configuration which could potentially resolve the issue. If you are still facing the issue, please create a new GitHub issue with your latest findings, with all the debugging information which could help us investigate.

Please follow the release notes to stay up to date with the latest developments which are happening in the Tensorflow space.

[github-actions[bot]]

This issue is stale because it has been open for 7 days with no activity. It will be closed if no further activity occurs. Thank you.

[github-actions[bot]]

This issue was closed because it has been inactive for 7 days since being marked as stale. Please reopen if you'd like to work on this further.