Need help on memory reduction

[sjiang23]

Hi Jax team,

The following is a highly simplified model of what I am really doing:

from jax import random, vmap, numpy as jnp

a = random.uniform(random.PRNGKey(0),(10**6,))
b = random.uniform(random.PRNGKey(1),(10**6,))

def foo(x, y):
    return jnp.sin(x + y)

batched_foo = vmap(foo, in_axes=(0, None), out_axes=0)
batched_foo(a, b)

The outcome is always memory leak:

UnfilteredStackTrace: RuntimeError: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 4000000000000 bytes.

The stack trace below excludes JAX-internal frames.
The preceding is the original exception that occurred, unmodified.

--------------------

The above exception was the direct cause of the following exception:

RuntimeError                              Traceback (most recent call last)
/usr/local/lib/python3.7/dist-packages/jax/_src/numpy/lax_numpy.py in deferring_binary_op(self, other)
   6551     if not isinstance(other, _accepted_binop_types):
   6552       return NotImplemented
-> 6553     return binary_op(self, other)
   6554   return deferring_binary_op
   6555 

RuntimeError: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 4000000000000 bytes.

The above operation is important in my own code and I believe there must be some workaround to this issue. However due to my limited knowledge on software engineering I am not able to find one. Could you please help me? Thank you!

[C-J-Cundy]

Hi, I don't think there's a simple plug-in method that will reduce the memory usage of the operation you've written.
The function you've written will return a 10^6 x 10^6-size matrix, i.e. one with a trillion entries. So just to store the output will require on the order of terabytes of memory.

If you're able to refactor your computational problem so that it doesn't need the entire matrix at once this would reduce the memory usage. So if e.g. you can deal with the rows one-by-one, you could loop through the rows with a jax.lax.scan. Inside the scan inner function you could vmap over however many rows you can fit in memory at once.

[sjiang23]

Thanks for your help, I will re-factor my problem.