title: Masking arbitrary square regions in batch in TensorFlow
body: <p>I have a large 3D tensor in TensorFlow 2.0 with shape [batch_size, N, N]. Each sample in the batch contains a square ROI that I want to mask out (the sizes of the ROIs are defined in a separate array named inner_dims). I've been googling quite a lot but haven't been able to find a solution that appears efficient.</p>
  
  <p><strong>Example:</strong></p>
  
  <p>The example below shows what I want to accomplish. In this case, I have manually created a mask that does what I want:</p>
  
  <pre class="lang-py prettyprint-override"><code>batch_size=5
  N=4
  # inputs.shape: [5, 4, 4]
  
  inner_dims=[2, 3, 1, 4, 2]
  
  
  mask=
  [
   [
    [1 1 0 0]
    [1 1 0 0]
    [0 0 0 0]
    [0 0 0 0]
   ]
   [
    [1 1 1 0]
    [1 1 1 0]
    [1 1 1 0]
    [0 0 0 0]
   ]
   [
    [1 0 0 0]
    [0 0 0 0]
    [0 0 0 0]
    [0 0 0 0]
   ]
   [
    [1 1 1 1]
    [1 1 1 1]
    [1 1 1 1]
    [1 1 1 1]
   ]
   [
    [1 1 0 0]
    [1 1 0 0]
    [0 0 0 0]
    [0 0 0 0]
   ]
  ]
  
  outputs=tf.multiply(inputs, mask)
  </code></pre>
  
  <p><strong>What I've tried</strong></p>
  
  <p>I tried the naïve approach of just assigning the affected values to zero, like so:</p>
  
  <pre class="lang-py prettyprint-override"><code>for i in range(0, len(inner_dims)):
      inputs[i, inner_dims[i]:, :] = 0
      inputs[i, :inner_dims[i], inner_dims[i]:] = 0
  </code></pre>
  
  <p>This, (as expected) gives me an item assignment error.</p>
  
  <p><code>TypeError: 'tensorflow.python.framework.ops.EagerTensor' object does not support item assignment</code></p>
  
  <p><strong>My question</strong></p>
  
  <p>How can I achieve the desired masking in a fast and GPU-friendly way? Is it better to generate a mask on the fly, or can the direct assignment approach be made to work?</p>