Adding tf.data optimization for rewriting `map(...).batch(...)` to `map_and_batch(...)`.