About jit of multi-host pmap not implemented

[matrix97317]

Q：NotImplementedError: jit of multi-host pmap not implemented (and jit-of-pmap can cause extra data movement anyway, so maybe you don't want it after all).

When I don't use jit, I run into out-of-memory issues, but when I use jit and a multi-machine training environment with pmap for all-reduce, I encounter compatibility problems between jit and pmap. Are there any good solutions