Repository: google/jax

Issue:
Nit: lax.conv_general_dilated does not support 0D-convolution
AFAIK 0-dimensional convolution should reduce to a fully-connected layer, but fails with the following example and error. While this is not a practically useful scenario, degenerate cases like these could be useful for testing.
```
from jax import lax
import jax.numpy as np

lax.conv_general_dilated(lhs=np.ones((10, 5)), 
                         rhs=np.ones((5, 7)), 
                         strides=(), 
                         paddign='SAME', 
                         dimension_numbers=('NC', 'IO', 'NC'))
```
```
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
<ipython-input-17-573e7695a148> in <module>()
----> 1 lax.conv_general_dilated(np.ones((10, 5)), np.ones((5, 7)), (), 'SAME', dimension_numbers=('NC', 'IO', 'NC'))

6 frames
google3/third_party/py/jax/lax/lax.py in conv_general_dilated(lhs, rhs, window_strides, padding, lhs_dilation, rhs_dilation, dimension_numbers, feature_group_count, precision)
    516       feature_group_count=feature_group_count,
    517       lhs_shape=lhs.shape, rhs_shape=rhs.shape,
--> 518       precision=_canonicalize_precision(precision))
    519 
    520 def dot(lhs, rhs, precision=None):

google3/third_party/py/jax/core.py in bind(self, *args, **kwargs)
    148     top_trace = find_top_trace(args)
    149     if top_trace is None:
--> 150       return self.impl(*args, **kwargs)
    151 
    152     tracers = map(top_trace.full_raise, args)

google3/third_party/py/jax/interpreters/xla.py in apply_primitive(prim, *args, **params)
    150 def apply_primitive(prim, *args, **params):
    151   """Impl rule that compiles and runs a single primitive 'prim' using XLA."""
--> 152   compiled_fun = xla_primitive_callable(prim, *map(arg_spec, args), **params)
    153   return compiled_fun(*args)
    154 

google3/third_party/py/jax/interpreters/xla.py in xla_primitive_callable(prim, *arg_specs, **params)
    174     device = device and next(d for d in all_devices if (type(d), d.id) == device)
    175   backend = xb.get_device_backend(device)
--> 176   aval_out = prim.abstract_eval(*avals, **params)
    177   if prim.multiple_results:
    178     handlers = tuple(map(aval_to_result_handler, aval_out))

google3/third_party/py/jax/lax/lax.py in standard_abstract_eval(prim, shape_rule, dtype_rule, *args, **kwargs)
   1500     return ConcreteArray(prim.impl(*[x.val for x in args], **kwargs))
   1501   elif least_specialized is ShapedArray:
-> 1502     return ShapedArray(shape_rule(*args, **kwargs), dtype_rule(*args, **kwargs))
   1503   elif least_specialized is UnshapedArray:
   1504     return UnshapedArray(dtype_rule(*args, **kwargs))

google3/third_party/py/jax/lax/lax.py in _conv_general_dilated_shape_rule(lhs, rhs, window_strides, padding, lhs_dilation, rhs_dilation, dimension_numbers, feature_group_count, **unused_kwargs)
   1943   lhs_trans = _dilate_shape(onp.take(lhs.shape, lhs_perm), lhs_dilation)
   1944   rhs_trans = _dilate_shape(onp.take(rhs.shape, rhs_perm), rhs_dilation)
-> 1945   out_trans = conv_shape_tuple(lhs_trans, rhs_trans, window_strides, padding)
   1946   return tuple(onp.take(out_trans, onp.argsort(out_perm)))
   1947 

google3/third_party/py/jax/lax/lax.py in conv_shape_tuple(lhs_shape, rhs_shape, strides, pads)
   4334     raise TypeError(msg.format(len(lhs_shape) - 2, len(pads)))
   4335 
-> 4336   lhs_padded = onp.add(lhs_shape[2:], onp.add(*zip(*pads)))
   4337   out_space = onp.floor_divide(
   4338       onp.subtract(lhs_padded, rhs_shape[2:]), strides) + 1

ValueError: invalid number of arguments
```


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
@@ -4320,7 +4320,8 @@ def conv_shape_tuple(lhs_shape, rhs_shape, strides, pads):
     msg = "Wrong number of explicit pads for convolution: expected {}, got {}."
     raise TypeError(msg.format(len(lhs_shape) - 2, len(pads)))
 
-  lhs_padded = onp.add(lhs_shape[2:], onp.add(*zip(*pads)))
+  lhs_padded = onp.add(lhs_shape[2:], onp.sum(onp.array(pads).reshape(-1, 2),
+                                              axis=1))
   out_space = onp.floor_divide(
       onp.subtract(lhs_padded, rhs_shape[2:]), strides) + 1
   out_space = onp.maximum(0, out_space)


Test patch:
diff --git a/tests/lax_test.py b/tests/lax_test.py
--- a/tests/lax_test.py
+++ b/tests/lax_test.py
@@ -489,6 +489,16 @@ def fun(lhs, rhs):
 
   # TODO(mattjj): test conv_general_dilated against numpy
 
+  def testConv0DIsDot(self):
+    rng = jtu.rand_default()
+    def args_maker():
+      return [rng((10, 5), onp.float32), rng((5, 7), onp.float32)]
+    jnp_fun = partial(lax.conv_general_dilated, window_strides=(),
+                      padding='VALID', dimension_numbers=('NC', 'IO', 'NC'))
+    self._CompileAndCheck(jnp_fun, args_maker, check_dtypes=True)
+    self._CheckAgainstNumpy(jnp_fun, onp.dot, args_maker)
+
+
   @staticmethod
   def _conv_transpose_via_grad(data, kernel, strides, padding,
                                rhs_dilation=None, dimension_numbers=None):