Estimate memory requirement from jaxpr

[joeryjoery]

Hi,

If I have some jax function f and I call e.g., jax.eval_shape(f) on this function for some input x, is there then also a way to get an estimate of the induced memory-requirements?

Could this be read from the generated jaxpr?

Even if this will result in an inaccurate estimation, it could be very useful for tuning batch-sizes programatically for calls to jax.vmap. Then I can opt for jax.lax.scan instead if my memory is insufficient, rather than crashing my program :(.

Any help is greatly appreciated :).

[jakevdp]

You can get a sense of the memory use using the cost_analysis() method of a lowered and compiled function. For example:

import jax
import jax.numpy as jnp

def f(x): 
  return jnp.sin(x) ** 2 + jnp.cos(x) ** 2

x = jnp.arange(1000)
compiled = jax.jit(f).lower(x).compile()
print(compiled.cost_analysis())

[{'transcendentals': 2000.0,
  'bytes accessed operand 0 {}': 4000.0,
  'bytes accessed output {}': 4000.0,
  'bytes accessed': 8000.0,
  'utilization operand 1 {}': 3.0,
  'bytes accessed operand 1 {}': 12000.0,
  'flops': 4000.0,
  'utilization operand 0 {}': 1.0}]

I have to admit, though, I don't have a great sense of how to interpret these outputs; there's some info here that may be helpful: https://jax.readthedocs.io/en/latest/aot.html#debug-information-and-analyses-when-available