how to apply mask with boolean_mask to the last dimensions

Viewed 494

I'm trying to apply a triangle mask to the last two dimensions of a tensor.

I'm doing something like this:

a, b, c, d = 2, 2, 2, 2
tensor = tf.random.uniform((a, b, c, d))

lower_triangular = tf.linalg.band_part(tf.reshape(tf.ones((c, d)), shape=(1, 1, c, d)), -1, 0)

tf.boolean_mask(tensor, lower_triangular == 0)

The last instruction fails with the following exception

/usr/local/lib/python3.6/dist-packages/tensorflow/python/util/dispatch.py in wrapper(*args, **kwargs)
    199     """Call target, and fall back on dispatchers if there is a TypeError."""
    200     try:
--> 201       return target(*args, **kwargs)
    202     except (TypeError, ValueError):
    203       # Note: convert_to_eager_tensor currently raises a ValueError, not a

/usr/local/lib/python3.6/dist-packages/tensorflow/python/ops/array_ops.py in boolean_mask_v2(tensor, mask, axis, name)
   1824   ```
   1825   """
-> 1826   return boolean_mask(tensor, mask, name, axis)
   1827 
   1828 

/usr/local/lib/python3.6/dist-packages/tensorflow/python/util/dispatch.py in wrapper(*args, **kwargs)
    199     """Call target, and fall back on dispatchers if there is a TypeError."""
    200     try:
--> 201       return target(*args, **kwargs)
    202     except (TypeError, ValueError):
    203       # Note: convert_to_eager_tensor currently raises a ValueError, not a

/usr/local/lib/python3.6/dist-packages/tensorflow/python/ops/array_ops.py in boolean_mask(tensor, mask, name, axis)
   1744     if axis_value is not None:
   1745       axis = axis_value
-> 1746       shape_tensor[axis:axis + ndims_mask].assert_is_compatible_with(shape_mask)
   1747 
   1748     leading_size = gen_math_ops.prod(shape(tensor)[axis:axis + ndims_mask], [0])

/usr/local/lib/python3.6/dist-packages/tensorflow/python/framework/tensor_shape.py in assert_is_compatible_with(self, other)
   1132     """
   1133     if not self.is_compatible_with(other):
-> 1134       raise ValueError("Shapes %s and %s are incompatible" % (self, other))
   1135 
   1136   def most_specific_compatible_shape(self, other):

ValueError: Shapes (2, 2, 2, 2) and (1, 1, 2, 2) are incompatible

If I create a mask with the exact same shape as my tensor, I got a strange result:

a, b, c, d = 2, 2, 2, 2
tensor = tf.random.uniform((a, b, c, d))

lower_triangular = tf.linalg.band_part(tf.ones((a, b, c, d)), -1, 0)

result = tf.boolean_mask(tensor, lower_triangular == 0)
result.shape

The result has a shape TensorShape([4]) not sure what is this doing.

What should be the right way of doing mask?

2 Answers

The result has a shape TensorShape([4]) not sure what is this doing.

You have a tensor like this:

   [[[[0.7786709 , 0.67790854],
     [0.72038054, 0.8475332 ]],

    [[0.945917  , 0.496732  ],
     [0.05510616, 0.44530773]]],


   [[[0.05596864, 0.59605396],
     [0.5459013 , 0.80567217]],

    [[0.1393286 , 0.74939907],
     [0.1695472 , 0.55127   ]]]]

and a mask like this:

  [[[[False,  True],
     [False, False]],

    [[False,  True],
     [False, False]]],


   [[[False,  True],
     [False, False]],

    [[False,  True],
     [False, False]]]]

So, you end up with a result like this:

[0.67790854, 0.496732  , 0.59605396, 0.74939907]

We can see that the mask "filters" out all values where the mask is false and flattens the result.

So if the total number of True values in your mask was n and you matched the dimensions of the tensor and mask you end up with a tensor of shape [n].

I am not clear what result you are expecting from your mask operation but, if the flattening surprised you, then maybe it sounds like you just want something like this?:

  [[[[0.        , 0.67790854],
     [0.        , 0.        ]],

    [[0.        , 0.496732  ],
     [0.        , 0.        ]]],


   [[[0.        , 0.59605396],
     [0.        , 0.        ]],

    [[0.        , 0.74939907],
     [0.        , 0.        ]]]]

which you can just get by multiplying by the mask (though, of course, we need to cast it first).

Maybe something like this:

tf.cast(lower_triangular == 0, tf.float32) * tensor

W.R.T the original error message:

If you want to apply a mask with the same rank as the tensor then the dimension all need to be equal. Whereas you can apply a mask of a lower number of dimensions (which is what I suspect you were trying to do).

For example:

tf.boolean_mask(tensor, lower_triangular[0,0] == 0)

gives:

[[[0.945917  , 0.496732  ],
  [0.05510616, 0.44530773]]]

We can also control which axis the mask applies to. So, for example:

tf.boolean_mask(tensor, lower_triangular[0,0] == 0, axis=2)

returns:

[[[0.67790854],
  [0.496732  ]],

  [[0.59605396],
  [0.74939907]]]

where the mask is applied to the 2nd axis (rather than the 0th axis as default)

tf.boolean_mask(
    tensor, mask, axis=None, name='boolean_mask'
)

Numpy equivalent is tensor[mask].

# 1-D example
tensor = [0, 1, 2, 3]
mask = np.array([True, False, True, False])
boolean_mask(tensor, mask)  # [0, 2]

In general, 0 < dim(mask) = K <= dim(tensor), and mask's shape must match the first K dimensions of tensor's shape. We then have: boolean_mask(tensor, mask)[i, j1,...,jd] = tensor[i1,...,iK,j1,...,jd] where (i1,...,iK) is the ith True entry of mask (row-major order). The axis could be used with mask to indicate the axis to mask from. In that case, axis + dim(mask) <= dim(tensor) and mask's shape must match the first axis + dim(mask) dimensions of tensor's shape.

For more information.

Related