tf.nest.map_structure