'Fixed point arithmetic with np.einsum

I'm trying to simulate fixed point arithmetic with np.einsum. For example, let's consider the following snippet that performs matrix multiplication:

A = np.random.uniform(size=(4, 5))
B = np.random.uniform(size=(5, 6))
C = np.einsum('ij,jk->ik', A, B)

We can compute the aforementioned example using fixed-point arithmetic as follows:

num_bits = 16
s = 2**num_bits

tmp = A.reshape((4, 5, 1)) * B.reshape((1, 5, 6))
tmp = np.round(tmp * s).astype(np.int32)
C = np.sum(tmp, axis=1) / s

The drawback of the aforementioned code is that the intermediate tensor tmp can be pretty huge since we're doing the outer product first and performing reduction after. Is there a more efficient way to implement it using numpy?



Sources

This article follows the attribution requirements of Stack Overflow and is licensed under CC BY-SA 3.0.

Source: Stack Overflow

Solution Source