tf.data filter dataset too slow
Hi @huangrt01 ,

I have tested the code with tf-nightly and found execution fails.Could you please check the [gist](https://colab.sandbox.google.com/gist/SuryanarayanaY/db40b44ba5857492008d9bcc7ab023d1/67330.ipynb) and confirm anything missing in submitted code? > Hi @huangrt01 ,
> 
> I have tested the code with tf-nightly and found execution fails.Could you please check the [gist](https://colab.sandbox.google.com/gist/SuryanarayanaY/db40b44ba5857492008d9bcc7ab023d1/67330.ipynb) and confirm anything missing in submitted code?

Hi @SuryanarayanaY ，
you can just ignore the error and the "batch+filter" benchmark. In fact, we have got the log results:
![image](https://github.com/tensorflow/tensorflow/assets/14976749/f82f9265-2a4e-4f58-ac3b-be89eb2222d4)


### Issue type

Performance

### Have you reproduced the bug with TensorFlow Nightly?

Yes

### Source

source

### TensorFlow version

2.4

### Custom code

Yes

### OS platform and distribution

_No response_

### Mobile device

_No response_

### Python version

_No response_

### Bazel version

_No response_

### GCC/compiler version

_No response_

### CUDA/cuDNN version

_No response_

### GPU model and memory

_No response_

### Current behavior?

Similar to issue #53169, I have observed that the "filter before batch" approach is significantly slow. Filtering the dataset alone takes 430ms, whereas the "batch+map" method only requires 20ms.
In theory, the computation of filter and map should be similar, but "filter before batch" consumes excessive time.
I attempted to filter after batching, but encountered a limitation where the filter predicate must return a scalar boolean value. Unfortunately, it does not support filtering batched elements.

My question is:
Is there a potential optimization for this performance issue? I aim to develop a customized operation that can filter batched elements (accepting [M,] shaped tensors as input and producing [N,] tensors as output). Is there a more efficient approach available?

```
import time

import tensorflow as tf
fast_dataset = tf.data.Dataset.range(10000)


def fast_benchmark(dataset, name, num_epochs=2):
    start_time = time.perf_counter()
    for _ in tf.data.Dataset.range(num_epochs):
        for _ in dataset:
            pass
    tf.print("Test", name, "Execution time(ms):", 1000 * (time.perf_counter() - start_time))


def increment(x):
    return x+1


def filter_fn(x):
  return tf.math.equal(tf.math.mod(x, 2), 1)


if __name__ == '__main__':
  fast_benchmark(
    fast_dataset
    .map(increment)
    .batch(256)
    ,
    "map+batch"
  )
  fast_benchmark(
    fast_dataset
    .batch(256)
    .map(increment)
    ,
    "batch+map"
  )
  fast_benchmark(
    fast_dataset
    .map(increment)
    .batch(256)
    .prefetch(tf.data.AUTOTUNE)
    ,
    "map+batch+prefetch"
  )
  fast_benchmark(
    fast_dataset
    .batch(256)
    .map(increment)
    .prefetch(tf.data.AUTOTUNE)
    ,
    "batch+map+prefetch"
  )
  fast_benchmark(
    fast_dataset
    .prefetch(tf.data.AUTOTUNE)
    .batch(256)
    .map(increment)
    ,
    "prefetch+batch+map"
  )
  fast_benchmark(
    fast_dataset
    .batch(256)
    .prefetch(tf.data.AUTOTUNE)
    .map(increment)
    ,
    "batch+prefetch+map"
  )
  fast_benchmark(
    fast_dataset
    .filter(filter_fn)
    .batch(256)
    ,
    "filter+batch"
  )
  fast_benchmark(
    fast_dataset
    .batch(256)
    .filter(filter_fn)
    ,
    "batch+filter"
  )
```

result:
![image](https://github.com/tensorflow/tensorflow/assets/14976749/0833e2b3-1850-412f-b9a8-92f701b307f9)


### Standalone code to reproduce the issue

```shell
import time

import tensorflow as tf
fast_dataset = tf.data.Dataset.range(10000)


def fast_benchmark(dataset, name, num_epochs=2):
    start_time = time.perf_counter()
    for _ in tf.data.Dataset.range(num_epochs):
        for _ in dataset:
            pass
    tf.print("Test", name, "Execution time(ms):", 1000 * (time.perf_counter() - start_time))


def increment(x):
    return x+1


def filter_fn(x):
  return tf.math.equal(tf.math.mod(x, 2), 1)


if __name__ == '__main__':
  fast_benchmark(
    fast_dataset
    .map(increment)
    .batch(256)
    ,
    "map+batch"
  )
  fast_benchmark(
    fast_dataset
    .batch(256)
    .map(increment)
    ,
    "batch+map"
  )
  fast_benchmark(
    fast_dataset
    .map(increment)
    .batch(256)
    .prefetch(tf.data.AUTOTUNE)
    ,
    "map+batch+prefetch"
  )
  fast_benchmark(
    fast_dataset
    .batch(256)
    .map(increment)
    .prefetch(tf.data.AUTOTUNE)
    ,
    "batch+map+prefetch"
  )
  fast_benchmark(
    fast_dataset
    .prefetch(tf.data.AUTOTUNE)
    .batch(256)
    .map(increment)
    ,
    "prefetch+batch+map"
  )
  fast_benchmark(
    fast_dataset
    .batch(256)
    .prefetch(tf.data.AUTOTUNE)
    .map(increment)
    ,
    "batch+prefetch+map"
  )
  fast_benchmark(
    fast_dataset
    .filter(filter_fn)
    .batch(256)
    ,
    "filter+batch"
  )
  fast_benchmark(
    fast_dataset
    .batch(256)
    .filter(filter_fn)
    ,
    "batch+filter"
  )
```


### Relevant log output

```shell
Test map+batch Execution time(ms): 585.0009880959988
Test batch+map Execution time(ms): 23.16068299114704
Test map+batch+prefetch Execution time(ms): 503.9997957646847
Test batch+map+prefetch Execution time(ms): 19.63987946510315
Test prefetch+batch+map Execution time(ms): 54.23441715538502
Test batch+prefetch+map Execution time(ms): 16.469698399305344
Test filter+batch Execution time(ms): 282.77427703142166
```
