Add tf.booleanMask() op.

[hiddentn]

TensorFlow.js version

0.11.1

Browser version

68

Describe the problem or feature request

i am currently trying to grap a few values from a tensor using Tensor.data() witch takes a lot of time given a large tensor.

const mask=tf.where(......)
const all_indices = tf.range(1,N+1,1"int32");
  //starting at 1 to prevent 0x0
const indices = mask.toInt().mul(all_indices)
const indicesArr = Array.from(await Indices.data()).filter(i => i > 0).map(i => i - 1)
const final = tf.gather(data, tf.tensor1d(indicesArr,"int32"))

i am the wondering when or if the python equivalent tf.boolean_mask will be implemented

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

Code to reproduce the bug / link to feature request

https://www.tensorflow.org/api_docs/python/tf/boolean_mask

[chenqing]

What is the current progress?
@dsmilkov

[ClaytonSmith]

This is a needed feature. booleanMaskAsync is an absolute pain to work with

[zqwitt]

Not sure if this will have any further discussion since this ticket has been closed. But is there any reason why this op needs to be asynchronous? I see it is using whereAsync() as the only internal async op but could this be implemented with where() instead, thus making is a synchronous op like the other gather and slicing ops and not return a Promise?

[mlajtos]

@zqwitt tf.where keeps the shape of the original tensor(s). tf.gather partially keeps the shape – it matches the provided indices tensor. Proposed tf.booleanMask dynamically changes the shape, this complicates things. But I agree that tf.booleanMaskAsync is a nightmare.

I found a way to do booleanMask sync with just one scalar read out (pseudo-code):

booleanMask(tensor, mask):
  count = sum(mask).arraySync() // sync scalar read out
  indices = range(0, tensor.shape[0])
  top = topk(where(mask, indices, -1), count)
  return gather(tensor, top)