I am looking for the equivalent of this code in pure tensorflow 2.1:
a_tensor[an_index_mask, 2] = 0
tf.boolean_mask that can reproduce
a_tensor[an_index_mask] but in my case
a_tensor is a 3-D matrix while
an_index_mask is only a 2-D mask.
I can’t really use numpy here as this occur in a
tf.data.Dataset.map and the tensors are not eager there.
In addition to this, I’d like to do some assignation on the extracted slice, but I may be able to solve this with some clever concatenation or other way.