Complete functionality of SparseTensor.__mul__

[vdhelm]

System information

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

Describe the feature and the current behavior/state.
Currently SparseTensor.mul is limited to multiplication between a sparse tensor and a dense tensor. Broadcasting is only supported from the dense side to the sparse side, as per the documentation: https://www.tensorflow.org/api_docs/python/tf/sparse/SparseTensor#__mul__

My request is to expand this behaviour to also include multiplications between two sparse tensors, and to allow broadcasting in both directions. This will complete the functionality and make the use of sparse tensors similar to that of dense tensors.

Will this change the current api? How?
The current options should of course still be supported, but more tensor combinations should be supported. This is purely the removal of a limitation that is already stated in the API.

Who will benefit with this feature?
Anyone who uses SparseTensors will want to manipulate them using the basic functions like multiplication and division (which can probably be added at the same time). This is certainly true for my own use in Graph Neural networks, but sparse tensor calculations are needed in many other applications as well.

Any Other info.
I have written a pure python implementation of broadcast multiply between two sparse tensors, which is shared in this stackoverflow question: https://stackoverflow.com/questions/63023958/how-to-efficiently-broadcast-multiply-two-sparse-tensors-in-tensorflow. However, as noted there, it is not memory efficient enough to be really useful. I expect that a C++ implementation is needed, and unfortunately I do not know how to write that. If I can contribute any more code, I would be happy to, but my knowledge of C++ and CUDA is limited.

[bhack]

We have also https://www.tensorflow.org/api_docs/python/tf/raw_ops/SparseMatrixSparseMatMul for SparseTensor (CSR).

[vdhelm]

We have also https://www.tensorflow.org/api_docs/python/tf/raw_ops/SparseMatrixSparseMatMul for SparseTensor (CSR).

Thank you, Yes, I have looked at that as well and it is very different from what I reference here. Note that matrix multiplication is a different operation than regular tensor multiplication with broadcasting. The fact that both are called multiplication is actually a little misleading, but that's the fault of mathematics terminology I guess.

[bhack]

You are right. I have re-read your ticket description and yes you was specifically requesting broadcasting in both directions.