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?