[tf.data] Optimize `Dataset.filter()` when the predicate returns one of its args.

This change avoids the overhead of function dispatch (~10--15us) when
the filter predicate simply returns one of its arguments directly. It
also adds a benchmark to track the performance of this optimization.

The checkpointing code required minor modifications to enable functions to be
instantiated in the `FilterDatasetOp::Compute()` method when an iterator is being
restored.

PiperOrigin-RevId: 188229570