Best practices for distributing memory?

[rkruegs123]

I have a memory-intensive computation that I want to scale up by using multiple GPUs. I have been experimenting with xmap, pmap, and shmap. My understanding is that these would typically be use to distribute batches in a training set to parallelize computation. Note that I am not training a neural network, and therefore the computation I am trying to parallelize is not analogous -- instead, the computation I am trying to vectorize (and distribute across GPUs) is in a scan.

I have found two things:

1. The only option for parallelizing across GPUs within a scan/jitted function is shmap. This is fine.
2. The memory does not distribute across GPUs. Instead, it is replicated. Therefore, this does not solve my issue.

So, does anybody know how to distribute memory across GPUs when the computation to be vectorized is within another function? A minimal example would be extremely helpful.

Thank you in advance.