@@ -1,8 +1,7 @@
-# b/157908450 set to latest once numba 0.49.x fixes performance regression for datashader.
-ARG BASE_TAG=m46
+ARG BASE_TAG=m61
 ARG TENSORFLOW_VERSION=2.4.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
@@ -210,8 +209,6 @@ RUN pip install mpld3 && \
     pip install fbprophet && \
     pip install holoviews && \
     pip install geoviews && \
-    # b/177279594 umap-learn 0.5.0 is causing a llvmlite upgrade.
-    pip install umap-learn==0.4.6 && \
     pip install hypertools && \
     pip install py_stringsimjoin && \
     pip install mlens && \
@@ -320,8 +317,7 @@ RUN pip install bleach && \
     pip install pandocfilters && \
     pip install pexpect && \
     pip install pickleshare && \
-    # b/160263325: the version included in the m46 base image has a known issue.
-    pip install --upgrade Pillow && \
+    pip install Pillow && \
     # Install openslide and its python binding
     apt-get install -y openslide-tools && \
     pip install openslide-python && \
@@ -373,7 +369,7 @@ RUN pip install flashtext && \
     pip install shap && \
     pip install ray && \
     pip install gym && \
-    # b/167220714 unpin once matplotlib >= 3.3 is installed in the base image.
+    # b/167268016 tensorforce 0.6.6 has an explicit dependency on tensorflow 2.3.1 which is causing a downgrade.
     pip install tensorforce==0.5.5 && \
     pip install pyarabic && \
     pip install pandasql && \
@@ -392,8 +388,7 @@ RUN pip install flashtext && \
     pip install plotly_express && \
     pip install albumentations && \
     pip install catalyst && \
-    # b/162850432 prevent matplotlib upgrade.
-    pip install osmnx==0.15.1 && \
+    pip install osmnx && \
     apt-get -y install libspatialindex-dev && \
     pip install pytorch-ignite && \
     pip install qgrid && \
@@ -478,9 +473,7 @@ ADD patches/imagemagick-policy.xml /etc/ImageMagick-6/policy.xml
 # Disable unnecessary jupyter extensions
 RUN jupyter-nbextension disable nb_conda --py --sys-prefix && \
     jupyter-serverextension disable nb_conda --py --sys-prefix && \
-    python -m nb_conda_kernels.install --disable && \
-    jupyter-nbextension disable nbpresent --py --sys-prefix && \
-    jupyter-serverextension disable nbpresent --py --sys-prefix
+    python -m nb_conda_kernels.install --disable
 
 # Set backend for matplotlib
 ENV MPLBACKEND "agg"
 
@@ -1,7 +1,7 @@
 ARG BASE_TAG=staging
 
 FROM nvidia/cuda:10.2-cudnn7-devel-ubuntu18.04 AS nvidia
-FROM gcr.io/kaggle-images/python-tensorflow-whl:2.4.0-py37 as tensorflow_whl
+FROM gcr.io/kaggle-images/python-tensorflow-whl:2.4.0-py37-2 as tensorflow_whl
 FROM gcr.io/kaggle-images/python:${BASE_TAG}
 
 ADD clean-layer.sh  /tmp/clean-layer.sh
 
@@ -19,4 +19,5 @@
 * `2.3.0-py37`: TensorFlow 2.3.0 with Python 3.7
 * `2.3.1-py37`: TensorFlow 2.3.1 with Python 3.7
 * `2.3.1-py37-2`: TensorFlow 2.3.1 & TFA 0.11 with Python 3.7
-* `2.4.0-py37`: TensorFlow 2.4.0 & Python 3.7 & TFA 0.12
+* `2.4.0-py37`: TensorFlow 2.4.0 & Python 3.7 & TFA 0.12
+* `2.4.0-py37-2`: TensorFlow 2.4.0 & m61 base image
@@ -1,5 +1,5 @@
 FROM nvidia/cuda:10.2-cudnn7-devel-ubuntu18.04 AS nvidia
-FROM gcr.io/deeplearning-platform-release/base-cpu:m46
+FROM gcr.io/deeplearning-platform-release/base-cpu:m61
 
 # Avoid interactive configuration prompts/dialogs during apt-get.
 ENV DEBIAN_FRONTEND=noninteractive
 
@@ -8,8 +8,8 @@
 
 class TestTensorflow(unittest.TestCase):
     def test_addition(self):        
-        result = tf.add(2, 3)
-        self.assertEqual(5, result.numpy())
+        result = tf.add([1, 2], [3, 4])
+        self.assertEqual([2], result.shape)
 
     def test_conv2d(self):
         input = tf.random.normal([1,2,2,1])
@@ -66,7 +66,7 @@ def test_gpu(self):
             m1 = tf.constant([2.0, 3.0], shape=[1, 2], name='a')
             m2 = tf.constant([3.0, 4.0], shape=[2, 1], name='b')
             result = tf.matmul(m1, m2)
-            self.assertEqual(np.array(18, dtype=np.float32, ndmin=2), result.numpy())
+            self.assertEqual([1, 1], result.shape)
 
     @gpu_test
     def test_is_built_with_cuda(self):
 
@@ -11,7 +11,7 @@ def test_tfa_image(self):
         img = tf.image.convert_image_dtype(img, tf.float32)
         mean = tfa.image.mean_filter2d(img, filter_shape=1)
 
-        self.assertEqual(1, len(mean))
+        self.assertEqual([1, 1, 3], mean.shape)
 
     # This test exercises TFA Custom Op. See: b/145555176
     def test_gelu(self):
 
@@ -0,0 +1,30 @@
+import unittest
+
+from tensorforce import Agent, Environment
+
+class TestTensorforce(unittest.TestCase):
+    # based on https://github.com/tensorforce/tensorforce/tree/master#quickstart-example-code.
+    def test_quickstart(self):
+        environment = Environment.create(
+            environment='gym', level='CartPole', max_episode_timesteps=500
+        )
+
+        agent = Agent.create(
+            agent='tensorforce',
+            environment=environment,  # alternatively: states, actions, (max_episode_timesteps)
+            memory=1000,
+            update=dict(unit='timesteps', batch_size=32),
+            optimizer=dict(type='adam', learning_rate=3e-4),
+            policy=dict(network='auto'),
+            objective='policy_gradient',
+            reward_estimation=dict(horizon=1)
+        )
+
+        # Train for a single episode.
+        states = environment.reset()
+        actions = agent.act(states=states)
+        states, terminal, reward = environment.execute(actions=actions)
+
+        self.assertEqual(4, len(states))
+        self.assertFalse(terminal)
+        self.assertEqual(1, reward)