Use custom TFA wheels.

This allows support for Custom Op. See https://github.com/tensorflow/addons#987

BUG=145555176

Author: rosbo

Dockerfile

@@ -1,7 +1,7 @@
 ARG BASE_TAG=latest
 ARG TENSORFLOW_VERSION=2.2.0
 
-FROM gcr.io/kaggle-images/python-tensorflow-whl:${TENSORFLOW_VERSION}-py37 as tensorflow_whl
+FROM gcr.io/kaggle-images/python-tensorflow-whl:${TENSORFLOW_VERSION}-py37-2 as tensorflow_whl
 FROM gcr.io/deeplearning-platform-release/base-cpu:${BASE_TAG}
 
 ADD clean-layer.sh  /tmp/clean-layer.sh
@@ -66,8 +66,9 @@ RUN pip install /tmp/tensorflow_gcs_config/tensorflow*.whl && \
     /tmp/clean-layer.sh
 
 # Install TensorFlow addons (TFA).
-# TFA functionnality relying on Custom Op won't work. See: https://github.com/tensorflow/addons/issues/987.
-RUN pip install tensorflow-addons==0.10.0 && \
+COPY --from=tensorflow_whl /tmp/tfa_cpu/*.whl /tmp/tfa_cpu/
+RUN pip install /tmp/tfa_cpu/tensorflow*.whl && \
+    rm -rf /tmp/tfa_cpu/ && \
     /tmp/clean-layer.sh
 
 RUN apt-get install -y libfreetype6-dev && \

gpu.Dockerfile

@@ -1,7 +1,7 @@
 ARG BASE_TAG=staging
 
 FROM nvidia/cuda:10.1-cudnn7-devel-ubuntu18.04 AS nvidia
-FROM gcr.io/kaggle-images/python-tensorflow-whl:2.2.0-py37 as tensorflow_whl
+FROM gcr.io/kaggle-images/python-tensorflow-whl:2.2.0-py37-2 as tensorflow_whl
 FROM gcr.io/kaggle-images/python:${BASE_TAG}
 
 ADD clean-layer.sh  /tmp/clean-layer.sh
@@ -102,6 +102,12 @@ RUN pip uninstall -y tensorflow && \
     pip install --no-deps mxnet-cu$CUDA_MAJOR_VERSION$CUDA_MINOR_VERSION && \
     /tmp/clean-layer.sh
 
+ # Reinstall TensorFlow addons (TFA) with GPU support.
+COPY --from=tensorflow_whl /tmp/tfa_gpu/*.whl /tmp/tfa_gpu/
+RUN pip install /tmp/tfa_gpu/tensorflow*.whl && \
+    rm -rf /tmp/tfa_gpu/ && \
+    /tmp/clean-layer.sh
+
 # Install GPU-only packages
 RUN pip install pycuda && \
     pip install cupy-cuda$CUDA_MAJOR_VERSION$CUDA_MINOR_VERSION && \