Does pmap with CPU backend broadcast it's in_axes=(None) items on CPU RAM?

[mohamad-amin]

I guess the question speaks for itself.

I'm trying to use jax.pmap to parallelize some part of my code on CPU, but I'm facing memory issues. According to my expectations, I should not run into any memory issue unless the broadcasted items in pmap are copied for each CPU.

P.S: I'm using XLA_FLAGS=--xla_force_host_platform_device_count=32 and my machine has 48 cores where each have 2 threads. So my batches are multiples of 32.