jaxlib build fixes. (#4066)

1. `wheel.pep425tags` has been removed as of
   pypa/setuptools#1829. Use the new
   `packaging.tags` instead.

2. Add `--allow-downgrades` to cuda install command. I'm not sure this
   is always necessary, but I ran into it, I'm guessing due to a cached
   docker image.

build/build_wheel_docker_entrypoint.sh

@@ -34,7 +34,7 @@ fi
 # Builds and activates a specific Python version.
 pyenv local "$PY_VERSION"
 
-PY_TAG=$(python -c "import wheel; import wheel.pep425tags as t; print(t.get_abbr_impl() + t.get_impl_ver())")
+PY_TAG=$(python -c "import packaging.tags as t; print(t.interpreter_name() + t.interpreter_version())")
 
 echo "Python tag: $PY_TAG"
 

build/install_cuda.sh

@@ -19,7 +19,7 @@ echo "cudnn version: $CUDNN_VERSION"
 
 apt-get update
 apt-get remove -y --allow-change-held-packages cuda-license-10-0 libcudnn7 libnccl2
-apt-get install -y --no-install-recommends \
+apt-get install -y --no-install-recommends --allow-downgrades \
   cuda-nvml-dev-$CUDA_VERSION \
   cuda-command-line-tools-$CUDA_VERSION \
   cuda-libraries-dev-$CUDA_VERSION \