How can I make numpy arrays the same shape before doing conditional computations?
04:36 03 Sep 2026

Let's say I have a conditional calculation like this

import numpy as np

def foo(a, b, c):
    if a <= 0 or c == 0:
        return foo_special(a, b, c)
    return foo_normal(a, b, c)

def foo_normal(a, b, c):
    return np.log(a) * b / c

def foo_special(a, b, c):
    return a + b + c

and I want to make foo work with arrays. A naive approach would be

def foo_naive(a, b, c):
    out = np.zeros_like(a)
    special = (a <= 0) | (c == 0)
    out[special] = foo_special(a[special], b[special], c[special])
    out[~special] = foo_normal(a[~special], b[~special], c[~special])
    return out

but this assumes that a, b and c have the same shape. However, normally, numpy allows me to mix scalars with arrays as well as arrays of different shapes, e.g.

foo_normal(
    4,
    np.array([1, 2, 3]),
    np.array([[5, 6, 7]]).T,
)

works just fine (returning a 3 x 3 array) and I want to keep that behavior.

Thus, what I need to do (I think), is to bring all input arrays into the same, final shape before giving them to foo_naive. How can I do this?

python numpy array-broadcasting numpy-ufunc