make_broadcastable.py

from einops import repeat, parse_shape


def make_broadcastable(array_list, pattern_list, result_pattern):
    shape_list = []
    all_keys = set()
    for array, pattern in zip(array_list, pattern_list):
        shape_list.append(parse_shape(array, pattern))
        all_keys |= set(shape_list[-1].keys())

    results = []
    for array, pattern, shape in zip(array_list, pattern_list, shape_list):
        shape.update({k: 1 for k in all_keys - set(shape.keys())})
        results.append(repeat(array, f'{pattern} -> {result_pattern}', **shape))

    return results


if __name__ == "__main__":
    import numpy as np

    xs = np.random.random(size=(2, 3))
    ys = np.random.random(size=(2, 5))
    zs = np.random.random(size=(2, 7))

    xs, ys, zs = make_broadcastable([xs, ys, zs], ['n x', 'n y', 'n z'], 'n x y z')

    print(xs.shape)  # (2, 3, 1, 1)
    print(ys.shape)  # (2, 1, 5, 1)
    print(zs.shape)  # (2, 1, 1, 7)