jax

7472ea3f - fix: use backend to call xb.process_count in _raise_warnings_or_errors_for_jit_of_pmap

fix: use backend to call xb.process_count in _raise_warnings_or_errors_for_jit_of_pmap not giving the backend to process_count was causing exception when lowering a functions for GPU precompilation device without having a GPU device.