def map_flat()

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

• 9 lines of code
• 6 McCabe index (conditional complexity)

def map_flat(f, values):
    """Apply the function f to 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]
    flat = tf.concat([tf.reshape(v, [-1]) for v in values], axis=0)
    flat = f(flat)
    parts = tf.split(flat, [tf.size(v) for v in values])
    return [tf.reshape(p, tf.shape(v)) for p, v in zip(parts, values)]