jax

8bdaf329 - fixed linting stuff for jax/_src/lax/control_flow files, added more skips to the multiprocess_gpu_test.py, reduced the skipped tests for test_dot, used test_util get_rocm_version function instead of duplicating it