message: Extend support of non-const axis tensor in boolean_mask
diff: @@ -1711,6 +1711,11 @@ def boolean_mask(tensor, mask, name="boolean_mask", axis=None):
  shape(tensor)[:axis], [leading_size],
  shape(tensor)[axis + ndims_mask:]
  ], 0))
  + # TODO(yongtang): tf.reshape in C++ kernel might have set the shape
  + # correctly, so the following may not be needed? It still might ben
  + # possible that there are some edge case where tensor_util.constant_value
  + # resolves more case than ShapeInference of tf.reshape in C++ kernel.
  + if axis_value is not None:
  first_dim = shape_tensor[axis:axis + ndims_mask].num_elements()
  tensor.set_shape(
  tensor_shape.as_shape(shape_tensor[:axis]).concatenate(
  