Question: I am new to JAX and trying to learn use it for running some code on a GPU. In my example I want to search for regular grids in a point cloud (for indexing X-ray diffraction data). With test_mats[4_000_000,3,3] the memory usage seems to be 15 MB. But with test_mats[5_000_000,3,3] I get an error about it wanting to allocate 19 GB. I can't tell whether this is a glitch in JAX, or because I am doing something wrong. My example code and output are below. I guess the problem is that it wants to create a temporary array of (N, 3, gvec.shape[1]) before doing the reduction, but I don't know how to see the memory profile for what happens inside the jitted/vmapped function. import sys import os import jax import jax.random import jax.profiler print('jax.version.__version__',jax.version.__version__) import scipy.spatial.transform import numpy as np # (3,N) integer grid spot positions hkls = np.mgrid[-3:4, -3:4, -3:4].reshape(3,-1) Umat = scipy.spatial.transform.Rotation.random( 10, random_state=42 ).as_matrix() a0 = 10.13 gvec = np.swapaxes( Umat.dot(hkls)/a0, 0, 1 ).reshape(3,-1) def count_indexed_peaks_hkl( ubi, gve, tol ): """ See how many gve this ubi can account for """ hkl_real = ubi.dot( gve ) hkl_int = jax.numpy.round( hkl_real ) drlv2 = ((hkl_real - hkl_int)**2).sum(axis=0) npks = jax.numpy.where( drlv2 < tol*tol, 1, 0 ).sum() return npks def testsize( N ): print("Testing size",N) jfunc = jax.vmap( jax.jit(count_indexed_peaks_hkl), in_axes=(0,None,None)) key = jax.random.PRNGKey(0) test_mats = jax.random.orthogonal(key, 3, (N,) )*a0 dev_gvec = jax.device_put( gvec ) scores = jfunc( test_mats, gvec, 0.01 ) jax.profiler.save_device_memory_profile(f"memory_{N}.prof") os.system(f"~/go/bin/pprof -top {sys.executable} memory_{N}.prof") testsize(400000) testsize(500000)  Output is: gpu4-03:~/Notebooks/JAXFits % python mem.py jax.version.__version__ 0.4.16 Testing size 400000 File: python Type: space Showing nodes accounting for 15.26MB, 99.44% of 15.35MB total Dropped 25 nodes (cum <= 0.08MB) flat flat% sum% cum cum% 15.26MB 99.44% 99.44% 15.26MB 99.44% __call__ 0 0% 99.44% 15.35MB 100% [python] 0 0% 99.44% 1.53MB 10.00% _pjit_batcher 0 0% 99.44% 15.30MB 99.70% _pjit_call_impl 0 0% 99.44% 15.30MB 99.70% _pjit_call_impl_python 0 0% 99.44% 15.30MB 99.70% _python_pjit_helper 0 0% 99.44% 15.35MB 100% bind 0 0% 99.44% 15.35MB 100% bind_with_trace 0 0% 99.44% 15.30MB 99.70% cache_miss 0 0% 99.44% 15.30MB 99.70% call_impl_cache_miss 0 0% 99.44% 1.53MB 10.00% call_wrapped 0 0% 99.44% 13.74MB 89.51% deferring_binary_op 0 0% 99.44% 15.35MB 100% process_primitive 0 0% 99.44% 15.30MB 99.70% reraise_with_filtered_traceback 0 0% 99.44% 15.35MB 100% testsize 0 0% 99.44% 1.53MB 10.00% vmap_f 0 0% 99.44% 15.31MB 99.74% wrapper Testing size 500000 2023-12-14 10:26:23.630474: W external/tsl/tsl/framework/bfc_allocator.cc:296] Allocator (GPU_0_bfc) ran out of memory trying to allocate 19.18GiB with freed_by_count=0. The caller indicates that this is not a failure, but this may mean that there could be performance gains if more memory were available. Traceback (most recent call last): File "~/Notebooks/JAXFits/mem.py", line 38, in <module> testsize(500000) File "~/Notebooks/JAXFits/mem.py", line 33, in testsize scores = jfunc( test_mats, gvec, 0.01 ) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ jaxlib.xla_extension.XlaRuntimeError: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 20596777216 bytes. -------------------- For simplicity, JAX has removed its internal frames from the traceback of the following exception. Set JAX_TRACEBACK_FILTERING=off to include these.``` 

Answer: The vmapped function is attempting to create an intermediate array of shape [N, 3, 3430]. For N=400_000, with float32 this amounts to 15GB, and for N=500_000 this amounts to 19GB. Your best option in this situation is probably to split your computation into sequentially-executed batches using lax.map or similar. Unfortunately there's not currently any automatic way to do that kind of chunked vmao, but there is a relevant feature request at https://github.com/google/jax/issues/11319, and there are some useful suggestions in that thread.