Numpy broadcasting implementation (for educational purposes)

broadcasting.py

import operator
from itertools import cycle
from numbers import Number
import numpy as np


def broadcast(a, b, op):
    result = []

    if isinstance(a, Number) and isinstance(b, Number):
        return op(a, b)

    if a.ndim == b.ndim:
        if a.shape[0] != b.shape[0]:
            if a.shape[0] == 1:
                a = cycle(a)
            elif b.shape[0] == 1:
                b = cycle(b)
            else:
                print(f'Could not broadcast together with shapes {a.shape} {b.shape}')
                return []
    elif a.ndim > b.ndim:
        b = cycle([b])
    else:
        a = cycle([a])

    for a_in, b_in in zip(a, b):
        result.append(broadcast(a_in, b_in, op))

    return np.array(result)


if __name__ == '__main__':
    list = []
    a = np.arange(9).reshape(3, 1, 3)
    b = np.arange(9).reshape(3, 3, 1)
    list += [(a, b)]

    a = np.arange(6).reshape(3, 2)
    b = np.arange(3)[:, None]
    list += [(a, b)]

    a = np.arange(6).reshape(2, 3)
    b = np.arange(3)
    list += [(a, b)]

    for a, b in list:
        result = broadcast(a, b, operator.add)
        np_result = a + b

        print(result)
        print('---------------------')
        print(np_result)
        print('****************************************************')

        np.testing.assert_equal(result, np_result)