How to use batch_vmap() in PR#19614

[green-slime]

Hi! Recently, I noticed that the batch_vmap() function has been submitted in PR #19614. In order to address the OOM issue with vmap(), I attempted to copy the submitted code into the JAX library and tried using it. However, while it returned the expected results in Example 1, it seems that Example 2 did not effectively demonstrate its ability to batch and avoid OOM problems. I believe I might have misunderstood how to use it. For instance, if I want to run 1e10 functions in 10 batches (let's say that the gpu can only handle 1e9 of them), how should I utilize the batch_vmap() function?
Thank you so much for any help!

@jit
def f(x):
    return jnp.array([1,2,3])
# Example 1:
print(vmap(f)(jnp.arange(3)))
print(batch_vmap(f,batch_size=2)(jnp.arange(3)))
# Example 2:
vmap(f)(jnp.arange(1e10)) # OOM
batch_vmap(f,batch_size=1e9)(jnp.arange(1e10)) # OOM

Example 1:

[[1 2 3]
 [1 2 3]
 [1 2 3]]
[[1 2 3]
 [1 2 3]
 [1 2 3]]

Example 2 (both of them):

2024-02-26 17:12:05.421499: W external/xla/xla/service/hlo_rematerialization.cc:2194] Can't reduce memory use below 23.81GiB (25567444992 bytes) by rematerialization; only reduced to 37.25GiB (40000000000 bytes)
2024-02-26 17:12:15.619517: W external/tsl/tsl/framework/bfc_allocator.cc:485] Allocator (GPU_0_bfc) ran out of memory trying to allocate 37.25GiB (rounded to 40000000000)requested by op 
2024-02-26 17:12:15.619963: W external/tsl/tsl/framework/bfc_allocator.cc:497] <allocator contains no memory>
2024-02-26 17:12:15.620116: E external/xla/xla/pjrt/pjrt_stream_executor_client.cc:2432] Execution of replica 0 failed: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 40000000000 bytes.
BufferAssignment OOM Debugging.

[jakevdp]

#19614 is an unmerged work-in-progress PR. I think the answer to "how to use it" is "don't yet" 😁

But if you want to weigh-in on the PR itself with the results of your testing, that may be helpful in making the new feature as useful as it can be when it eventually lands!

[green-slime]

Please accept my apologies. I have come to realize a glaring mistake on my part: it was the array itself exceeding the cache limit, and not related to vmap. I am truly embarrassed by my oversight. :(