title: tf.data filter dataset using label predicate
body: <p>I am trying to filter the CIFAR10 training and test data with specific labels as given below,</p>
  
  <pre><code>import tensorflow as tf
  from tensorflow.keras import datasets, layers, models
  import tensorflow_datasets as tfds
  import matplotlib.pyplot as plt
  import numpy as np
  
  </code></pre>
  
  <h3>Dataset</h3>
  
  <pre><code>dataset = datasets.cifar10.load_data()
  
  </code></pre>
  
  <h3>Split the dataset</h3>
  
  <pre><code>train_data = tf.data.Dataset.from_tensor_slices((dataset[0][0],dataset[0][1]))
  test_data = tf.data.Dataset.from_tensor_slices((dataset[1][0],dataset[1][1]))
  </code></pre>
  
  <h3>Filter function</h3>
  
  <pre><code>def filter_f(datas,filter_labels = tf.constant([0,1,2])):
    x = tf.not_equal(datas[1],filter_labels)
    x = tf.reduce_sum(tf.cast(x, tf.uint8))
    return tf.greater(x, tf.constant(0,tf.uint8))
  
  dataset = train_data.filter(filter_f).batch(200)
  
  </code></pre>
  
  <p>as per <a href="https://stackoverflow.com/questions/55731774/filter-dataset-to-get-just-images-from-specific-class">similar issue</a>. However, the filter function returns the unfiltered in the above code. </p>
  
  <pre><code>labels = []
  for i, x in enumerate(tfds.as_numpy(dataset)):
      labels.append(x[1][0][0])
  print(labels)
  </code></pre>
  
  <p>Returns</p>
  
  <pre><code>[4, 7, 5, 6, 0, 5, 5, 6, 5, 3, 6, 7, 0, 0, 6, 3]
  </code></pre>
  
  <p>To reproduce the result, please use this <a href="https://colab.research.google.com/drive/1IJuaQlUhDfA_SrA1ldtpvpo-q4j_fvWp" rel="nofollow noreferrer">colab link</a></p>