jax.random seems to have unnecessary buffer allocations on stack about jax

[py4]

@froystig nice! it worked

[jakevdp]

Have you tried it under jit?

[py4]

@jakevdp yeah same result:

import jax
import jax.numpy as jnp

@jax.jit
def get_random_array(key):
    return jax.random.uniform(key, (524288, 4096), dtype=jnp.bfloat16, minval=-1, maxval=1)

if __name__ == "__main__":
    key = jax.random.key(0)
    jax.profiler.start_trace("/tmp/profile/tmp")
    arr = get_random_array(key)
    arr.block_until_ready()
    jax.profiler.stop_trace()

[froystig]

What RNG implementation are you using? If it's the default (threefry), have you tried using it in "partitionable" mode? You can do that by adding the line:

jax.config.update('jax_threefry_partitionable', True)

at the top of your file.

#18480 has more details (and other ways to set the option).

I ask because partitionable Threefry mode often involves fewer reshapes and stacking.

[py4]

@froystig the original issue is resolved, however FYI, the context manager seems not to be working (the memory issue apears with context manager, but not with jax.config.update('jax_threefry_partitionable', True) at top:

import jax
import jax.numpy as jnp

@jax.jit
def get_random_array(key):
    with jax.threefry_partitionable(True):
        return jax.random.uniform(key, (524288, 4096), dtype=jnp.bfloat16, minval=-1, maxval=1)

if __name__ == "__main__":
    key = jax.random.key(0)
    jax.profiler.start_trace("/tmp/profile/tmp")
    arr = get_random_array(key)
    arr.block_until_ready()
    jax.profiler.stop_trace()

[froystig]

Thanks! Minimized and filed as #21061.