<p>I have a partial (not complete because I don't thoroughly understand your goals) answer found here: <a href="https://arxiv.org/pdf/1108.1320.pdf" rel="nofollow noreferrer">Compressed matrix multiplication</a>.</p>

<p>You are probably asking for the problem known as "Matrix Multiplication with Sparse Output".</p>

<p>Read the first section (introduction): </p>

<blockquote>
  <p>"Our method can be seen as a compressed sensing method for the matrix product, with the nonstandard idea that the <strong>SKETCH</strong> of AB is computed <strong>WITHOUT EXPLICITLY CONSTRUCTING</strong> AB"</p>
</blockquote>

<p>Under some assumptions (sparseness of the output matrix and number of sketches) and using <strong>ERROR CORRECTING CODES</strong> one can obtain the largest entries of the initial product without multiplying A and B.</p>

<p>Look at figure 2 for an example.</p>
matrix multiplication with sparse output specified by another sparse matrix on tensorflow
<p>I would like to implement a matrix multiplication on TensorFlow like <em>C</em> = <em>A</em> · <em>B</em> where <em>A</em> ∈ ℝ<sup><em>n</em>,<em>k</em></sup> and <em>B</em> ∈ ℝ<sup><em>k</em>,<em>n</em></sup>. In my case, <em>n</em> could be potentially large but <em>k</em> is typically small (e.g. low rank or latent embeddings).</p>

<p>As you know the dense matrix <em>C</em> ∈ ℝ<sup><em>n</em>,<em>n</em></sup> is expensive to store in RAM. However, the only entries in <em>C</em> I want to retain are sparse. That is, by defining another sparse matrix <em>D</em> ∈ ℝ<sup><em>n</em>,<em>n</em></sup>, what I really care about are those values with indices <code>[i,j]</code> having values in <em>D</em>. The non-empty values in <em>D</em> should only be <code>1</code>.</p>

<p>So, instead of doing stuff like this:</p>

<pre><code>tmp = tf.matmul(A,B)
C = tf.SparseTensor(D.indices, tf.gather_nd(tmp, D.indices)*D.values, D.dense_shape)
</code></pre>

<p>I want to avoid explicitly computing the dense tensor <code>tmp</code> above.</p>

<p>Thanks in advance!</p>