New dtype style testing

Bihaqo · Oct 30, 2018 · b690957 · b690957
1 parent fea76f6
commit b690957
Showing 1 changed file with 17 additions and 9 deletions.
diff --git a/t3f/autodiff_test.py b/t3f/autodiff_test.py
@@ -8,13 +8,13 @@
 from t3f import autodiff
 
 
-class AutodiffTest(tf.test.TestCase):
+class _AutodiffTest():
 
   def testGradients(self):
-    w = initializers.random_matrix(([5] * 3, None))
-    A = initializers.random_matrix(([5] * 3, [5] * 3))
-    x = initializers.random_matrix(([5] * 3, None))
-    z = initializers.random_matrix(([5] * 3, None))
+    w = initializers.random_matrix(([5] * 3, None), dtype=self.dtype)
+    A = initializers.random_matrix(([5] * 3, [5] * 3), dtype=self.dtype)
+    x = initializers.random_matrix(([5] * 3, None), dtype=self.dtype)
+    z = initializers.random_matrix(([5] * 3, None), dtype=self.dtype)
 
     def func1(x):
       return 0.5 * ops.flat_inner(x, w) ** 2
@@ -36,10 +36,10 @@ def func2(x):
       np.testing.assert_allclose(actual_v2, desired_v2, rtol=1e-4)
 
   def testHessianVectorProduct(self):
-    w = initializers.random_matrix(([5] * 3, None))
-    A = initializers.random_matrix(([5] * 3, [5] * 3))
-    x = initializers.random_matrix(([5] * 3, None))
-    z = initializers.random_matrix(([5] * 3, None))
+    w = initializers.random_matrix(([5] * 3, None), dtype=self.dtype)
+    A = initializers.random_matrix(([5] * 3, [5] * 3), dtype=self.dtype)
+    x = initializers.random_matrix(([5] * 3, None), dtype=self.dtype)
+    z = initializers.random_matrix(([5] * 3, None), dtype=self.dtype)
     projected_vector = ops.full(riemannian.project(z, x))
 
     def func1(x):
@@ -68,6 +68,14 @@ def func2(x):
       np.testing.assert_allclose(actual2_v, desired2_v, rtol=1e-3)
 
 
+class AutodiffTestFloat32(tf.test.TestCase, _AutodiffTest):
+  dtype = tf.float32
+
+
+class AutodiffTestFloat64(tf.test.TestCase, _AutodiffTest):
+  dtype = tf.float64
+
+
 if __name__ == "__main__":
   tf.test.main()