pythonnumpymatrixvectorizationnumba# Fastest way to add matrices of different shapes in Python/Numba

I want to "add" two matrices, a matrix `a`

with shape (`K,T`

) and a matrix `b`

of shape `(K,N)`

, to result in a matrix of shape `(K,T,N`

)

The following works ok:

```
import numpy as np
from numba import njit
@njit
def add_matrices(a, b):
K, T, N = a.shape[0], a.shape[1], b.shape[1]
result_matrix = np.empty((K, T, N))
for k in range(K):
for t in range(T):
for n in range(N):
result_matrix[k, t, n] = a[k, t] + b[k, n]
return result_matrix
K = 10
T = 11
N = 12
a = np.ones((K,T))
b = np.ones((K,N))
result = add_matrices(a, b)
```

Is there a faster (vectorized?) way to do it that doesn't require the for loops, which I think is slowing down the function, especially for larger values of `K, T,N`

?

Solution

Use broadcasting.

```
a[:,:,None] + b[:,None,:]
```

This makes `a`

appear to have size [K, T, 1], and `b`

to have size `[K, 1, N]`

. Numpy knows how to add these two together.

- Python Jinja2 LaTeX Table
- Getting attributes of a class
- How can I print many significant figures in Python?
- How to allow list append() method to return the new list
- Calculate Last Friday of Month in Pandas
- Python type hint for Iterable[str] that isn't str
- How to iterate over a list in chunks
- How to exit the entire application from a Python thread?
- Running shell command and capturing the output
- How do I pass a variable by reference?
- Convert range(r) to list of strings of length 2 in python
- How can I get the start and end dates for each week?
- how to use send_message() in python-telegram-bot
- Python conditional replacement based on element type
- How can I count the number of items in an arbitrary iterable (such as a generator)?
- Find longest consecutive range of numbers in list
- Insert text in braces with asyncpg
- How does one put a link / url to the web-site's home page in Django?
- How to determine if a path is a subdirectory of another?
- Custom Keybindings for Ipython terminal
- FastAPI asynchronous background tasks blocks other requests?
- How to make sure that information from one file is duplicated into several text documents, without specific lines
- Installing a Python environment with Anaconda
- sklearn pipeline model predicting same results for all input
- Brew command not found after installing Anaconda Python
- How to get an XPath from selenium webelement or from lxml?
- Pipe PuTTY console to Python script
- How to align the axes of a figure in matplotlib?
- Persist ParentDocumentRetriever of langchain
- How to reset index in a pandas dataframe?