Install blas & mkl (#1060)

* Install blas & mkl

These are necessary for the TPU pytorch XLA.
They also provide a significant perforance boost.

They used to be automatically included when we were installing torch
from conda but they are not now that we install torch from pip.

I added a test to prevent regression.

http://b/195414519

* separate conda commands to speed up resolution

Author: rosbo

Dockerfile

@@ -42,6 +42,7 @@ RUN conda config --add channels conda-forge && \
     conda config --add channels nvidia && \
     conda config --add channels rapidsai && \
     # ^ rapidsai is the highest priority channel, default lowest, conda-forge 2nd lowest.
+    conda install mkl blas && \
     # b/161473620#comment7 pin required to prevent resolver from picking pysal 1.x., pysal 2.2.x is also downloading data on import.
     conda install cartopy=0.19 imagemagick=7.0 pyproj==3.1.0 pysal==2.1.0 && \
     /tmp/clean-layer.sh

tests/test_numpy.py

@@ -12,5 +12,5 @@ def test_array(self):
     # Numpy must be linked to the MKL. (Occasionally, a third-party package will muck up the installation
     # and numpy will be reinstalled with an OpenBLAS backing.)
     def test_mkl(self):
-        # This will throw an exception if the MKL is not linked correctly.
-        get_info("blas_mkl")
+        # This will throw an exception if the MKL is not linked correctly or return an empty dict.
+        self.assertTrue(get_info("blas_mkl"))