def test_threaded_map_structure_fn_with_error(self):
fn = lambda x: x / 0
result = array_utils.threaded_map_structure(fn, STRUCT)
tf.nest.assert_same_structure(result, STRUCT)
expected_flat = [None] * len(tf.nest.flatten(STRUCT))
self.assertEqual(tf.nest.flatten(result), expected_flat)
def test_threaded_map_structure_2args(self):
fn = lambda x, y: x + y
actual = array_utils.threaded_map_structure(fn, STRUCT, STRUCT)
expected = tf.nest.map_structure(fn, STRUCT, STRUCT)
self.assertEqual(actual, expected)
