def map_flat_chunked()

in lm_human_preferences/utils/core.py [0:0]

• 8 lines of code
• 5 McCabe index (conditional complexity)

def map_flat_chunked(f, values, *, limit=1<<29):
    """
    Apply the function f to chunked, flattened, concatenated values, then split and reshape back to original shapes.
    """
    values = tuple(values)
    for v in values:
        assert not isinstance(v, tf.IndexedSlices)
    values = [tf.convert_to_tensor(v) for v in values]
    chunks = chunk_tensors(values, limit=limit)
    mapped_values = [v for chunk in chunks for v in map_flat(f, chunk)]
    return mapped_values