Fix issue in boolean_mask when axis is passed as a tensor
This PR tries to address the issue raised in 32236 where
a TypeError was thrown out when axis is passed as a tensor.
In the docstring axis has been specified as accepting a 1-D tensor.
This PR fixes the issue.
This PR fixes 32236.
Signed-off-by: Yong Tang <yong.tang.github@outlook.com>
diff --git a/tensorflow/python/ops/array_ops.py b/tensorflow/python/ops/array_ops.py
index fbb977f..6374e87 100644
--- a/tensorflow/python/ops/array_ops.py
+++ b/tensorflow/python/ops/array_ops.py
@@ -1699,6 +1699,9 @@
"Number of mask dimensions must be specified, even if some dimensions"
" are None. E.g. shape=[None] is ok, but shape=None is not.")
axis = 0 if axis is None else axis
+ axis_value = tensor_util.constant_value(axis)
+ if axis_value is not None:
+ axis = axis_value
shape_tensor[axis:axis + ndims_mask].assert_is_compatible_with(shape_mask)
leading_size = gen_math_ops.prod(shape(tensor)[axis:axis + ndims_mask], [0])
