title: tensorflow consumes too much memory when using tf.map_fn
body: <p>I have the following code:</p>
  <pre><code>    @tf.function
      def c(self, point_instances, instances, alpha):
          def inner_comp(j):
              print(j.shape)
              print(tf.transpose(instances).shape)
              print(alpha.shape)
              print(tf.linalg.matvec(tf.transpose(instances), alpha).shape)
              tmp = tf.tensordot(j, tf.linalg.matvec(tf.transpose(instances), alpha), 1)
              print(tmp.shape)
              return tmp
  
          return tf.reduce_max(tf.abs(tf.map_fn(inner_comp, point_instances)))
  </code></pre>
  <p>I noticed that when <code>tf.map_fn</code> finishes the second iteration, it gets killed due to memory shortage. So the console output would be like:</p>
  <pre><code>(245245,)
  (245245, 460)
  (460,)
  (245245,)
  ()
  
  (245245,)
  (245245, 1040)
  (1040,)
  (245245,)
  ()
  
  
  Killed
  </code></pre>
  <p>Although my data is a large-scale, the output of <code>inner_comp</code> is always a scalar. I tried to allocate 200 GB of memory but still it gets killed. What would be the reason?</p>