Mathematica equivalent of NumPy broadcast addition for large arrays

Question

What is the Mathematica equivalent of the following Python code with the vectors' broadcast addition?

import numpy as np

a = np.random.rand(5000, 1, 5);
b = np.random.rand(1, 500, 5);

result = a + b #shape: (5000, 500, 5) 

Answer 1

For your specific case (dimension 1 only in the first two slots), this might work:

a = RandomReal[{0, 1}, {5000, 1, 5}];
b = RandomReal[{0, 1}, {1, 500, 5}];
c1 = Flatten[
 Outer[Plus, a, b, 2],
 {{1, 3}, {2, 4}}
 ]// RepeatedTiming // First

0.255

It is a bit more tedious to use Compile but also a bit faster:

Creating the CompiledFunctions:

cf = Compile[{{a, _Real, 2}},
   Table[Flatten[a, 1], {500}],
   CompilationTarget -> "WVM",
   RuntimeAttributes -> {Listable},
   Parallelization -> True
   ];
cg = Compile[{{b, _Real, 3}},
   Table[Flatten[b, 1], {5000}],
   CompilationTarget -> "WVM",
   RuntimeAttributes -> {Listable},
   Parallelization -> True
   ];

Running the actual code:

c2 = Plus[cf[a], cg[b]]; // RepeatedTiming // First
Max[Abs[c1 - c2]]

0.19

0.

Final remarks

The general case may be treated by a suitable combination of ArrayReshape, MapThread, Outer, and Flatten. Or, maybe even better, by ad-hoc compliled, Listable CompiledFunctions such as cf and cg instead of MapThread. Anyways, one would probably need a thourough case analysis for that.

Answer 2

You can use:

a = RandomReal[{0,1}, {3,5}]
b = RandomReal[{0,1}, {8,5}]
c = Table[a[[i]] + b[[j]], {i, Length[a]}, {j, Length[b]}]

Dimension[c] will be {3,8,5}, similar to result.shape in Python of (3,8,5)