Search code examples
pythonnumpytensorflowcrashkernel

Python: Alternatives to numpy or tensorflow tensordot for large matrices


I am looking for more efficient solutions to calculating a tensor product between large matrices in Python. Specifically, I want to calculate a tensor product such as $H \otimes I$ for two matrices H and I. However, for large matrices of dimension 156 and larger my kernel crashes. I am quite new to coding, so there likely could be an obvious solution I'm missing.

I have created a reproducible example:

import numpy as np
import tensorflow as tf
A = np.random.random((200,200))
B = np.random.random((200,200))
C = tf.tensordot(A,B,axes=0)

Tensorflow gives a warning 2024-05-07 10:57:10.174373: W external/local_tsl/tsl/framework/cpu_allocator_impl.cc:83] Allocation of 12800000000 exceeds 10% of free system memory. Indicating that I need to find a more memory-efficient way of calculating this. Numpy does not give a similar warning, but I assume it is the same underlying problem. Doing this with much smaller matrices do not cause the same problems.

What alternatives do I have to calculate large tensor products of matrices?


Solution

  • According to the docs, when a and b are matrices (order 2), the case axes=0 gives the outer product, a tensor of order 4.

    So the size of the result is 200 * 200 * 200 * 200 floats (12.8 GB, as Tensorflow warns). This is not a problem of memory efficiency, just the size of the requested operation.