tf.keras.backend.map_fn