Repository: google/jax

Issue:
conv_general_dilated do not support complex number
I am currently shifting a project from autograd to Jax which involve the convolution for complex number. In autograd, the convolution function from autograd/autograd/scipy/signal.py/convolve do not restrict the type of input, it works pretty well. However, neither jax/jax/scipy/signal.py/convolve nor jax/jax/lax/lax.py/conv_general_dilated support the usage of data format, complex64. Would i know if there's any plan for supporting complex64 or any direction for me to make it works like HIPS/autograd library? Thank you so much. 


Discussion:
Can you share a bit more information? What did you run, and on what hardware platform?
`---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
<ipython-input-6-761948a0f0be> in <module>
      5                                (1,1),  # lhs/image dilation
      6                                (1,1),  # rhs/kernel dilation
----> 7                                dn)     # dimension_numbers = lhs, rhs, out dimension permutation
      8 print("out shape: ", out.shape)
      9 print("First output channel:")

~/.local/lib/python3.6/site-packages/jax/lax/lax.py in conv_general_dilated(lhs, rhs, window_strides, padding, lhs_dilation, rhs_dilation, dimension_numbers, feature_group_count, batch_group_count, precision)
    540       batch_group_count=batch_group_count,
    541       lhs_shape=lhs.shape, rhs_shape=rhs.shape,
--> 542       precision=_canonicalize_precision(precision))
    543 
    544 def dot(lhs: Array, rhs: Array, precision: Optional[PrecisionType] = None) -> Array:

~/.local/lib/python3.6/site-packages/jax/core.py in bind(self, *args, **kwargs)
    197     top_trace = find_top_trace(args)
    198     if top_trace is None:
--> 199       return self.impl(*args, **kwargs)
    200 
    201     tracers = map(top_trace.full_raise, args)

~/.local/lib/python3.6/site-packages/jax/interpreters/xla.py in apply_primitive(prim, *args, **params)
    164 def apply_primitive(prim, *args, **params):
    165   """Impl rule that compiles and runs a single primitive 'prim' using XLA."""
--> 166   compiled_fun = xla_primitive_callable(prim, *map(arg_spec, args), **params)
    167   return compiled_fun(*args)
    168 

~/.local/lib/python3.6/site-packages/jax/interpreters/xla.py in xla_primitive_callable(prim, *arg_specs, **params)
    172   device = _device_from_arg_devices(arg_devices)
    173   backend = xb.get_device_backend(device)
--> 174   aval_out = prim.abstract_eval(*avals, **params)
    175   if not prim.multiple_results:
    176     handle_result = aval_to_result_handler(device, aval_out)

~/.local/lib/python3.6/site-packages/jax/lax/lax.py in standard_abstract_eval(prim, shape_rule, dtype_rule, *args, **kwargs)
   1697     return ConcreteArray(prim.impl(*[x.val for x in args], **kwargs))
   1698   elif least_specialized is ShapedArray:
-> 1699     return ShapedArray(shape_rule(*args, **kwargs), dtype_rule(*args, **kwargs))
   1700   elif least_specialized is UnshapedArray:
   1701     return UnshapedArray(dtype_rule(*args, **kwargs))

~/.local/lib/python3.6/site-packages/jax/lax/lax.py in _conv_general_dilated_dtype_rule(lhs, rhs, window_strides, padding, lhs_dilation, rhs_dilation, dimension_numbers, **unused_kwargs)
   2230     dimension_numbers, **unused_kwargs):
   2231   return naryop_dtype_rule(_input_dtype, [_float, _float],
-> 2232                           'conv_general_dilated', lhs, rhs)
   2233 
   2234 _conv_spec_transpose = lambda spec: (spec[1], spec[0]) + spec[2:]

~/.local/lib/python3.6/site-packages/jax/lax/lax.py in naryop_dtype_rule(result_dtype, accepted_dtypes, name, *avals, **kwargs)
   1737       typename = str(onp.dtype(aval_dtype).name)
   1738       typenames = ', '.join(t.__name__ for t in types)
-> 1739       raise TypeError(msg.format(name, typename, i, i, typenames))
   1740   _check_same_dtypes(name, False, *aval_dtypes)
   1741   return result_dtype(*avals)

TypeError: conv_general_dilated does not accept dtype complex64 at position 0. Accepted dtypes at position 0 are subtypes of floating.`

I tried to follow the guide from Advanced Jax tutorials - the sharp bits - convolution part and amend it with complex value, the above error is what i get after running the code. From my understanding, the code from lax.py
`def _conv_general_dilated_dtype_rule(
    lhs, rhs, *, window_strides, padding, lhs_dilation, rhs_dilation,
    dimension_numbers, **unused_kwargs):
  return naryop_dtype_rule(_input_dtype, [_float, _float],
                          'conv_general_dilated', lhs, rhs)`
restricted the input type to float but not complex number, would changing this function  (_conv_general_dilated_dtype_rule) lead to the support of complex-valued convolution? Thank you so much.
When i make the amendment changing the data type 
`2020-07-10 16:07:06.243313: F external/org_tensorflow/tensorflow/compiler/xla/service/gpu/gpu_conv_runner.cc:491] %custom-call = (c64[1,200,198,3]{2,1,3,0}, u8[0]{0}) custom-call(c64[1,200,198,3]{2,1,3,0} %copy.3, c64[3,3,3,3]{1,0,2,3} %copy.4), window={size=3x3 pad=1_1x1_1}, dim_labels=b01f_01io->b01f, custom_call_target="__cudnn$convForward", metadata={op_type="conv_general_dilated" op_name="conv_general_dilated[ batch_group_count=1\n                      dimension_numbers=ConvDimensionNumbers(lhs_spec=(0, 3, 1, 2), rhs_spec=(3, 2, 0, 1), out_spec=(0, 3, 1, 2))\n                      feature_group_count=1\n                      lhs_dilation=(1, 1)\n                      lhs_shape=(1, 200, 198, 3)\n                      padding=((1, 1), (1, 1))\n                      precision=None\n                      rhs_dilation=(1, 1)\n                      rhs_shape=(3, 3, 3, 3)\n                      window_strides=(1, 1) ]"}, backend_config="{\"algorithm\":\"0\",\"tensor_ops_enabled\":false,\"conv_result_scale\":1,\"activation_mode\":\"0\",\"side_input_scale\":0}"`

And here is the program i run for complex convolution

> import numpy as onp
import jax
import jax.numpy as jnp
from jax import lax
from jax.lax import conv_general_dilated
from matplotlib import pyplot as plt
\
img = onp.zeros((1, 200, 198, 3), dtype=jnp.complex64)
for k in range(3):
    x = 30 + 60*k
    y = 20 + 60*k
    img[0, x:x+10, y:y+10, k] = 1.0
print('image shape',img.shape)
\
kernel = onp.zeros((3, 3, 3, 3), dtype=jnp.complex64)
kernel += onp.array([[1, 1, 0],
                [1, 0,-1],
                [0,-1,-1]])[:, :, onp.newaxis, onp.newaxis]
print('kernel shape:',kernel.shape)
\
dn = lax.conv_dimension_numbers(img.shape,     # only ndim matters, not shape
                                kernel.shape,  # only ndim matters, not shape
                                ('NHWC', 'HWIO', 'NHWC'))  # the important bit
print(dn)
\
out = lax.conv_general_dilated(img,    # lhs = image tensor
                               kernel, # rhs = conv kernel tensor
                               (1,1),  # window strides
                               'SAME', # padding mode
                               (1,1),  # lhs/image dilation
                               (1,1),  # rhs/kernel dilation
                               dn)     # dimension_numbers = lhs, rhs, out dimension permutation
print("out shape: ", out.shape)

Is it the problem related to XLA compiler?

Resolving patch:
diff --git a/jax/lax/lax.py b/jax/lax/lax.py
--- a/jax/lax/lax.py
+++ b/jax/lax/lax.py
@@ -2402,7 +2402,7 @@ 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 naryop_dtype_rule(_input_dtype, [_float, _float],
+  return naryop_dtype_rule(_input_dtype, [_float | _complex, _float | _complex],
                           'conv_general_dilated', lhs, rhs)
 
 _conv_spec_transpose = lambda spec: (spec[1], spec[0]) + spec[2:]
@@ -2499,16 +2499,37 @@ def _conv_general_dilated_transpose_rhs(
       feature_group_count=feature_group_count,
       batch_group_count=batch_group_count, precision=precision)
 
+
 def _conv_general_dilated_translation_rule(
-    c, lhs, rhs, *, window_strides, padding, lhs_dilation, rhs_dilation,
-    dimension_numbers, feature_group_count, batch_group_count, precision,
-    **unused_kwargs):
+    c, lhs, rhs, *, window_strides, padding,
+    lhs_dilation, rhs_dilation, dimension_numbers, feature_group_count,
+    batch_group_count, precision, expand_complex_convolutions, **unused_kwargs):
   assert type(dimension_numbers) is ConvDimensionNumbers
   dimension_numbers = _conv_general_proto(dimension_numbers)
-  return xops.ConvGeneralDilated(lhs, rhs, window_strides, padding, lhs_dilation,
-                                 rhs_dilation, dimension_numbers,
-                                 feature_group_count, batch_group_count,
-                                 precision_config=_precision_config(precision))
+  precision_config = _precision_config(precision)
+  dtype = c.get_shape(lhs).numpy_dtype()
+  conv = lambda x, y: xops.ConvGeneralDilated(
+      x, y, window_strides, padding, lhs_dilation, rhs_dilation,
+      dimension_numbers, feature_group_count, batch_group_count,
+      precision_config=precision_config)
+  if expand_complex_convolutions and onp.issubdtype(dtype, onp.complexfloating):
+    # We use a trick for complex multiplication due to Gauss which uses three
+    # multiplications and five additions; instead of the naive method of four
+    # multiplications and two additions.
+    # https://en.wikipedia.org/wiki/Multiplication_algorithm#Complex_multiplication_algorithm
+    #
+    # This performance win comes with a trade-off in accuracy; especially in
+    # cases when the real and imaginary differ hugely in magnitude. The relative
+    # error bound (e.g. 1p-24 in case of float32) would be relative to the
+    # maximum of real and imaginary parts of the result instead of being
+    # satisfied by the real and imaginary parts independently of each other.
+    lhs_real, lhs_imag = xops.Real(lhs), xops.Imag(lhs)
+    rhs_real, rhs_imag = xops.Real(rhs), xops.Imag(rhs)
+    k1 = conv(xops.Add(lhs_real, lhs_imag), rhs_real)
+    k2 = conv(lhs_real, xops.Sub(rhs_imag, rhs_real))
+    k3 = conv(lhs_imag, xops.Add(rhs_real, rhs_imag))
+    return xops.Complex(xops.Sub(k1, k3), xops.Add(k1, k2))
+  return conv(lhs, rhs)
 
 def _conv_general_dilated_batch_rule(
     batched_args, batch_dims, *, window_strides, padding,
@@ -2631,7 +2652,16 @@ def _conv_general_dilated_masking_rule(
 
 conv_general_dilated_p = standard_primitive(
     _conv_general_dilated_shape_rule, _conv_general_dilated_dtype_rule,
-    'conv_general_dilated', _conv_general_dilated_translation_rule)
+    'conv_general_dilated', partial(_conv_general_dilated_translation_rule,
+                                    expand_complex_convolutions=False))
+
+# TODO(b/161124619, b/161126248): XLA does not support complex convolution on
+# CPU or GPU; on these backends, lower complex convolutions away.
+xla.backend_specific_translations['cpu'][conv_general_dilated_p] = partial(
+    _conv_general_dilated_translation_rule, expand_complex_convolutions=True)
+xla.backend_specific_translations['gpu'][conv_general_dilated_p] = partial(
+    _conv_general_dilated_translation_rule, expand_complex_convolutions=True)
+
 ad.defbilinear(conv_general_dilated_p,
                _conv_general_dilated_transpose_lhs,
                _conv_general_dilated_transpose_rhs)


Test patch:
diff --git a/tests/lax_autodiff_test.py b/tests/lax_autodiff_test.py
--- a/tests/lax_autodiff_test.py
+++ b/tests/lax_autodiff_test.py
@@ -346,7 +346,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 grad_float_dtypes
+      for dtype in grad_inexact_dtypes
       for padding in ([((0, 0), (0, 0)), ((1, 0), (0, 1))] +
         ([((0, -1), (0, 0))] if lhs_shape[2] != 0 else []))
       for dim_nums, perms in [
diff --git a/tests/lax_test.py b/tests/lax_test.py
--- a/tests/lax_test.py
+++ b/tests/lax_test.py
@@ -463,7 +463,7 @@ def numpy_fun(lhs, rhs):
            (j * feature_group_count * batch_group_count, i, 4, 5))
           for w in [0, 10]
           for b, i, j in itertools.product([2, 3], repeat=3)]
-      for dtype in float_dtypes for strides in [(1, 1), (2, 1)]
+      for dtype in inexact_dtypes for strides in [(1, 1), (2, 1)]
       for padding in [((1, 2), (2, 0)), ((10, 8), (7, 13))]
       for lhs_dilation, rhs_dilation in itertools.product(
           [(1, 1), (1, 2), (1, 4)], repeat=2)