title: Tensorflow gather columns of a matrix very slow
body: <p>Given two matrices A (1000 x 100) and B (100 x 1000), instead of directly computing their product in tensorflow, i.e., <code>tf.dot(A,B)</code>, I want to first select 10 cols (randomly) from A and 10 rows from B and then use the <code>tf.dot(A_s,B_s)</code></p>
  
  <p>Naturally, the second multiplication should be much faster as the number of required multiplications reduces by factor of 10.</p>
  
  <p>However, in reality, it seems selecting given columns of matrix A in tensorflow to creat A_s is an extremly inefficient process.  </p>
  
  <p>Given indices of the required columns in <code>idx</code>, I tried the following solutions to creat A_s. The solutions are ranked according to their performance:</p>
  
  <ol>
  <li><code>.   A_s = tf.transpose(tf.gather(tf.unstack(A, axis=1), idx))</code>: </li>
  </ol>
  
  <p><code>tf.dot(A_s,B_s)</code> 5 times slower than <code>tf.dot(A,B)</code> because creating A_s is too expensive.</p>
  
  <ol start="2">
   2. 
  </ol>
  
  <pre><code>
  
       p_shape = K.shape(params)
       p_flat = K.reshape(params, [-1])
       i_flat = K.reshape(K.reshape(
          K.arange(0, p_shape[0]) * p_shape[1], [-1, 1]) + indices, [-1])
       indices = [i_flat]
       v = K.transpose(indices)
       updates = i_flat * 0 - 1
       shape = tf.to_int32([p_shape[0] * p_shape[1]])
       scatter = tf.scatter_nd(v, updates, shape) + 1
       out_temp = tf.dynamic_partition(p_flat,
                       partitions=scatter, num_partitions=2)[0]
       A_s = tf.reshape(out_temp, [p_shape[0], -1])
  
  </code></pre>
  
  <p>results in 6-7 times slower product</p>
  
  <ol start="3">
   3. 
  </ol>
  
  <pre><code>
  
        X,Y =  tf.meshgrid((tf.range(0, p_shape[0])),indices)
        idx = K.concatenate([K.expand_dims(
             K.reshape((X),[-1]),1), 
             K.expand_dims(K.reshape((Y),[-1]),1)],axis=1)
        A_s = tf.reshape(tf.gather_nd(params, idx), [p_shape[0], -1])
  
  </code></pre>
  
  <p>10-12 times slower.</p>
  
  <p>Any idea on how I can improve the efficiency of column selection process is very much appreciated.</p>
  
  <p>PS1: I ran all the experiments on CPU.</p>
  
  <p>PS2: Matrix A is a placeholder not a variable. In some implementation it can get problematic as its shape may not be inferred.</p>