Pythonic way to perform broadcasting binary operations between two Ndarrays

Question

What the title says. I am trying to do the following tasks:

1) perform a binary operation(like +, -, *, /, >, <) between two 2-d numpy arrays, A1, A2, where A1.shape = (N1, N2_1), A2.shape = (N1, N2_2) to generate a 3-d numpy array whose shape is (N, N2_1, N2_2)

2) perform a binary operation(like +, -, *, /, >, <) between two 3-d numpy arrays, A1, A2, where A1.shape = (N1, N2, N3_1), A2.shape = (N1, N2, N3_2) to generate a 4-d numpy array whose shape is (N, N2, N3_1, N3_2)

I find myself writing very unpythonic code to do the first task 1). I will appreciate if someone can show me the proper pythonic code example to get it done.

---

my attempt:

sources of data:

import numpy as np
n_row = 10000
n_col_a1 = 3
n_col_a2 = 4

a1 = np.tile(np.arange(n_col_a1), (n_row, 1))
a2 = np.tile(np.arange(n_col_a2), (n_row, 1))

my unpythonic broadcasting attempt:

X1 = np.broadcast_to(a1, (n_col_a2, *a1.shape))
X1 = np.moveaxis(X1, 0, -1)

X2 = np.broadcast_to(a2, (n_col_a1, *a2.shape))
X2 = np.moveaxis(X2, 0, -2)

result_a1_minus_a2 = X1 - X2
print(result_a1_minus_a2)

Answer 1

The two input arrays have to be expanded to:

(N1, N2_1)     => (N1, N2_1, 1)
(N1, N2_2)     => (N1, 1,    N2_2)
(N1, N2_1, N2_2)

e.g.

A1[:, :, None] * A2[:, None, :]

Similarly:

(N1, N2, N3_1)
(N1, N2, N3_2)
(N1, N2, N3_1, N3_2)

A1[...,None] * A2[:,:,None,:]