Repository: google/jax

Issue:
lax.conv does not support float64
The lax package currently does not support float64 convolutions. This appears to be an explicit constraint and not a bug.

https://github.com/google/jax/blob/master/jax/lax/lax.py#L1828

Feature request: allow convolutions with float64 -> float64 in full precision. TF supports this type of convolution, so I don't think there's anything inherent preventing this, but maybe it's something subtle.


Discussion:

Resolving patch:
diff --git a/jax/lax/lax.py b/jax/lax/lax.py
--- a/jax/lax/lax.py
+++ b/jax/lax/lax.py
@@ -1504,7 +1504,6 @@ def _brcast_to(x, shape):
     return broadcast(x, shape)
 
 
-_f32 = {onp.float32}
 _float = {onp.floating}
 _complex = {onp.complexfloating}
 _complex_elem_types = {onp.float32, onp.float64}
@@ -1825,8 +1824,8 @@ def _conv_general_dilated_shape_rule(
 def _conv_general_dilated_dtype_rule(
     lhs, rhs, window_strides, padding, lhs_dilation, rhs_dilation,
     dimension_numbers, **unused_kwargs):
-  return binop_dtype_rule(_input_dtype, [_f32, _f32], 'conv_general_dilated',
-                          lhs, rhs)
+  return binop_dtype_rule(_input_dtype, [_float, _float],
+                          'conv_general_dilated', lhs, rhs)
 
 _conv_spec_transpose = lambda spec: (spec[1], spec[0]) + spec[2:]
 _conv_sdims = lambda spec: spec[2:]


Test patch:
diff --git a/tests/lax_test.py b/tests/lax_test.py
--- a/tests/lax_test.py
+++ b/tests/lax_test.py
@@ -309,7 +309,7 @@ def testConcatenateAgainstNumpy(self, dim, base_shape, dtype, num_arrs, rng):
       for lhs_shape, rhs_shape in [
           ((b, i, 9, 10), (j, i, 4, 5))
           for b, i, j in itertools.product([2, 3], repeat=3)]
-      for dtype in [onp.float32]
+      for dtype in float_dtypes
       for strides in [(1, 1), (1, 2), (2, 1)]
       for padding in ["VALID", "SAME"]
       for rng in [jtu.rand_small()]))
@@ -331,7 +331,7 @@ def fun(lhs, rhs):
       for lhs_shape, rhs_shape in [
           ((b, i, 9, 10), (j, i, 4, 5))
           for b, i, j in itertools.product([2, 3], repeat=3)]
-      for dtype in [onp.float32]
+      for dtype in float_dtypes
       for strides in [(1, 1), (1, 2), (2, 1)]
       for padding in ["VALID", "SAME"]
       for rng in [jtu.rand_small()]))
@@ -354,7 +354,8 @@ def testConvAgainstNumpy(self, lhs_shape, rhs_shape, dtype, strides, padding,
       for lhs_shape, rhs_shape in [
           ((b, i, 9, 10), (j, i, 4, 5))
           for b, i, j in itertools.product([1, 2, 3], repeat=3)]
-      for dtype in [onp.float32] for strides in [(1, 1), (1, 2), (2, 1)]
+      for dtype in float_dtypes
+      for strides in [(1, 1), (1, 2), (2, 1)]
       for padding in [((0, 0), (0, 0)), ((1, 2), (2, 0))]
       for lhs_dilation, rhs_dilation in itertools.product(
           [(1, 1), (1, 2), (2, 2)], repeat=2)
@@ -418,7 +419,7 @@ def numpy_fun(lhs, rhs):
       for lhs_shape, rhs_shape in [
           ((b, i, 9, 10), (j, i, 4, 5))
           for b, i, j in itertools.product([2, 3], repeat=3)]
-      for dtype in [onp.float32] for strides in [(1, 1), (2, 1)]
+      for dtype in float_dtypes for strides in [(1, 1), (2, 1)]
       for padding in [((1, 2), (2, 0))]
       for lhs_dilation, rhs_dilation in itertools.product(
           [(1, 1), (1, 2)], repeat=2)
@@ -494,7 +495,7 @@ def _transpose_conv_kernel(data, kernel, dimension_numbers):
       for lhs_shape, rhs_shape in [
           ((b, 9, 10, i), (k, k, j, i))  # NB: i,j flipped in RHS for transpose
           for b, i, j, k in itertools.product([2,3],[2,3],[2,3],[3,4,5])]
-      for dtype in [onp.float32]
+      for dtype in float_dtypes
       for strides in [(1, 1), (1, 2), (2, 1), (2, 2), (3, 3)]
       for padding in ["VALID", "SAME"]
       for dspec in [('NHWC', 'HWIO', 'NHWC'),]
@@ -527,7 +528,7 @@ def fun_via_grad(lhs, rhs):
       for lhs_shape, rhs_shape in [
           ((b, 9, 10, i), (k, k, i, j))
           for b, i, j, k in itertools.product([2,3],[2,3],[2,3],[3,4,5])]
-      for dtype in [onp.float32]
+      for dtype in float_dtypes
       for strides in [(1, 1), (1, 2), (2, 1), (2, 2), (3, 3)]
       for padding in ["VALID", "SAME"]
       for dspec in [('NHWC', 'HWIO', 'NHWC'),]
@@ -559,7 +560,7 @@ def fun_via_grad(lhs, rhs):
       for lhs_shape, rhs_shape in [
           ((b, 10, i), (k, i, j))
           for b, i, j, k in itertools.product([2,3],[2,3],[2,3],[3,4,5])]
-      for dtype in [onp.float32]
+      for dtype in float_dtypes
       for strides in [(1,), (2,), (3,)]
       for padding in ["VALID", "SAME"]
       for dspec in [('NHC', 'HIO', 'NHC'),]
@@ -1646,7 +1647,7 @@ def testConcatenateGrad(self, dim, base_shape, dtype, num_arrs, rng):
             for b, i, j in itertools.product([2, 3], repeat=3)],
            [((4, 2, 1), (3, 2, 1), [(1,)])])
        for strides in all_strides
-       for dtype in [onp.float32]
+       for dtype in float_dtypes
        for padding in ["VALID", "SAME"]
        for rng in [jtu.rand_small()]))
   @jtu.skip_on_devices("tpu")
@@ -1678,7 +1679,7 @@ def testConvGrad(self, lhs_shape, rhs_shape, dtype, strides, padding, rng):
        for strides in all_strides
        for rhs_dil in rhs_dils
        for lhs_dil in lhs_dils
-       for dtype in [onp.float32]
+       for dtype in float_dtypes
        for padding in all_pads
        for rng in [jtu.rand_small()]))
   @jtu.skip_on_devices("tpu")
@@ -1715,7 +1716,7 @@ def testConvWithGeneralPaddingGrad(self, lhs_shape, rhs_shape, dtype, strides,
       for strides in all_strides
       for rhs_dil in rhs_dils
       for lhs_dil in lhs_dils
-      for dtype in [onp.float32]
+      for dtype in float_dtypes
       for padding in all_pads
       for dim_nums, perms in [
           (("NCHW", "OIHW", "NCHW"), ([0, 1, 2, 3], [0, 1, 2, 3])),