Repository: google-deepmind/dm-haiku

Issue:
Automatically test float64 compatibility on CPU and GPU
We should add the equivalent of this test:
https://github.com/deepmind/dm-haiku/blob/4bcbec9f874c847075ad666c7f3621c8c524af3b/haiku/_src/integration/bfloat16_test.py#L31

...except using float64. This involves enabling `--jax_enable_x64` for the test & fixing outstanding errors.


Discussion:

Resolving patch:
diff --git a/haiku/_src/base.py b/haiku/_src/base.py
--- a/haiku/_src/base.py
+++ b/haiku/_src/base.py
@@ -443,7 +443,7 @@ def transform(f, *, apply_rng=False) -> Transformed:
       out = apply(params, rng, *a, **k)
 
   Note that the ``rng`` argument is typically not required for `apply` and
-  passing ``None`` is accpeted.
+  passing ``None`` is accepted.
 
   The first thing to do is to define a `Module`. A module encapsulates some
   parameters and a computation on those parameters:


Test patch:
diff --git a/haiku/_src/integration/bfloat16_test.py b/haiku/_src/integration/bfloat16_test.py
--- a/haiku/_src/integration/bfloat16_test.py
+++ b/haiku/_src/integration/bfloat16_test.py
@@ -25,19 +25,19 @@
 import jax.numpy as jnp
 
 
-class Bfloat16Test(parameterized.TestCase):
-
-  @test_utils.combined_named_parameters(descriptors.ALL_MODULES)
-  def test_bfloat16(
+class DTypeTestCase(parameterized.TestCase):
+  def assert_dtype(
       self,
+      test_dtype: DType,
       module_fn: descriptors.ModuleFn,
       shape: Shape,
-      dtype: DType,
+      input_dtype: DType,
   ):
+    """ Checks that modules accepting float32 input_dtype output test_dtype. """
     if jax.local_devices()[0].platform != 'tpu':
       self.skipTest('bfloat16 only supported on TPU')
 
-    if dtype != jnp.float32:
+    if input_dtype != jnp.float32:
       self.skipTest('Skipping module without float32 input')
 
     rng = jax.random.PRNGKey(42)
@@ -54,23 +54,35 @@ def g(x):
     x = jax.random.uniform(rng, shape)
     params, state = init_fn(rng, x)
 
-    # Cast f32 to bf16.
-    f32_to_bf16 = (
-        lambda v: v.astype(jnp.bfloat16) if v.dtype == jnp.float32 else v)
-    params, state = jax.tree_map(f32_to_bf16, (params, state))
+    # Cast f32 to test_dtype.
+    f32_to_test_dtype = (
+        lambda v: v.astype(test_dtype) if v.dtype == jnp.float32 else v)
+    params, state = jax.tree_map(f32_to_test_dtype, (params, state))
 
-    # bf16 in should result in bf16 out.
-    x = x.astype(jnp.bfloat16)
+    # test_dtype in should result in test_dtype out.
+    x = x.astype(test_dtype)
 
     for _ in range(2):
       y, state = apply_fn(params, state, rng, x)
 
-      def assert_bf16(v):
+      def assert_dtype(v):
         if v.dtype != jnp.int32:
-          self.assertEqual(v.dtype, jnp.bfloat16)
+          self.assertEqual(v.dtype, test_dtype)
+
+      jax.tree_map(assert_dtype, y)
+      jax.tree_map(assert_dtype, state)
+
+
+class Bfloat16Test(DTypeTestCase):
+  @test_utils.combined_named_parameters(descriptors.ALL_MODULES)
+  def test_bfloat16(
+      self,
+      module_fn: descriptors.ModuleFn,
+      shape: Shape,
+      dtype: DType,
+  ):
+    self.assert_dtype(jnp.bfloat16, module_fn, shape, dtype)
 
-      jax.tree_map(assert_bf16, y)
-      jax.tree_map(assert_bf16, state)
 
 if __name__ == '__main__':
   absltest.main()
diff --git a/haiku/_src/integration/float64_test.py b/haiku/_src/integration/float64_test.py
new file mode 100644
--- /dev/null
+++ b/haiku/_src/integration/float64_test.py
@@ -0,0 +1,41 @@
+# Lint as: python3
+# Copyright 2020 DeepMind Technologies Limited. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Tests for haiku._src.conformance.float64_test."""
+# This is put into a separate file since it requires enabling
+# float64 right after loading of jax.
+
+from absl.testing import absltest
+from haiku._src import test_utils
+from haiku._src.integration import descriptors
+from haiku._src.integration.bfloat16_test import DTypeTestCase
+from haiku._src.typing import DType, Shape  # pylint: disable=g-multiple-import
+from jax.config import config
+import jax.numpy as jnp
+config.update("jax_enable_x64", True)
+
+
+class Float64Test(DTypeTestCase):
+  @test_utils.combined_named_parameters(descriptors.ALL_MODULES)
+  def test_float64(
+      self,
+      module_fn: descriptors.ModuleFn,
+      shape: Shape,
+      dtype: DType,
+  ):
+    self.assert_dtype(jnp.float64, module_fn, shape, dtype)
+
+if __name__ == '__main__':
+  absltest.main()