Add FLAX (#1105)

Includes a smoke test.

http://b/207406362

Author: rosbo

Dockerfile.tmpl

@@ -137,6 +137,7 @@ RUN pip install pysal && \
     pip install seaborn python-dateutil dask python-igraph && \
     pip install pyyaml joblib husl geopy ml_metrics mne pyshp && \
     pip install pandas && \
+    pip install flax && \
     # Install h2o from source.
     # Use `conda install -c h2oai h2o` once Python 3.7 version is released to conda.
     apt-get install -y default-jre-headless && \

tests/test_flax.py

@@ -0,0 +1,15 @@
+import unittest
+
+import jax.numpy as jnp
+import numpy as np
+
+from flax import linen as nn
+
+
+class TestFlax(unittest.TestCase):
+
+    def test_bla(self):
+        x = jnp.full((1, 3, 3, 1), 2.)
+        mul_reduce = lambda x, y: x * y
+        y = nn.pooling.pool(x, 1., mul_reduce, (2, 2), (1, 1), 'VALID')
+        np.testing.assert_allclose(y, np.full((1, 2, 2, 1), 2. ** 4))