#tf-data Add a range check for `index_flat_map`.
#tf-data Add a range check for `index_flat_map`.

Without this check, it could crash if `index_map_fn` returns an
out-of-bound index due to:
https://github.com/tensorflow/tensorflow/blob/69a908420c5c5b90027f23905cd842c76ca3955c/tensorflow/core/framework/tensor.cc#L1104
