Repository: google/jax
Issue: MAML Gradients + Convnets don't work

The attached python code snippet implements a MAML gradient update on 1) an MLP and 2) a CNN.

When MODEL_TYPE = 'mlp', everything works. When MODEL_TYPE = 'conv', I get the following error:

RuntimeError: Invalid argument: Expected LHS feature dimension (value 1) to be a multiple of feature_group_count (value 1), and LHS feature dimension / feature_group_count = RHS feature dimension (value 40); got <conv>(f32[1,1,40,40], f32[10,1,1,40])
Dimension numbers: {kernel_input_feature_dimension: 3
kernel_spatial_dimensions: 1
kernel_spatial_dimensions: 2
input_batch_dimension: 3
output_batch_dimension: 2
output_feature_dimension: 3
input_spatial_dimensions: 1
input_spatial_dimensions: 2
output_spatial_dimensions: 0
[maml-cnn.txt](https://github.com/google/jax/files/2867892/maml-cnn.txt)

output_spatial_dimensions: 1
}.:     @     0x7ff2befc9214  xla::XlaBuilder::ReportErrorOrReturn()
    @     0x7ff2befd6eaa  xla::ConvGeneralDilated()
    @     0x7ff2b82e4cb8  xla::swig::LocalComputationBuilder::ConvGeneralDilated()
    @     0x7ff029565901  _wrap_LocalComputationBuilder_ConvGeneralDilated()
    @           0x4bc454  PyEval_EvalFrameEx

Note that for both model architectures, the maml objective loss can be computed. It's only in the case of attempting to take a gradient of the CNN-based loss, things break.



Changed files:
jax/lax.py, tests/lax_test.py

diff --git a/jax/lax.py b/jax/lax.py
index 2c6c62b13..e37a20c91 100644
--- a/jax/lax.py
+++ b/jax/lax.py
@@ -1136,7 +1136,7 @@ def _conv_general_dilated_transpose_lhs(
   lhs_sdims, rhs_sdims, out_sdims = map(_conv_sdims, dimension_numbers)
   lhs_spec, rhs_spec, out_spec = dimension_numbers
   t_rhs_spec = _conv_transpose(rhs_spec)
-  trans_dimension_numbers = ConvDimensionNumbers(lhs_spec, t_rhs_spec, out_spec)
+  trans_dimension_numbers = ConvDimensionNumbers(out_spec, t_rhs_spec, lhs_spec)
   padding = _conv_general_vjp_lhs_padding(
       onp.take(lhs_shape, lhs_sdims), onp.take(rhs_shape, rhs_sdims),
       window_strides, onp.take(g.shape, out_sdims), padding, lhs_dilation,
@@ -3194,7 +3194,10 @@ def remaining(original, *removed_lists):
   blacklist = set(itertools.chain(*removed_lists))
   return [i for i in original if i not in blacklist]
 
-
+# lhs_spec and out_spec are lists containing
+#   [batch dim, feature dim, spatial dims ...]
+# rhs_spec is a list containing:
+#   [out feature dim, in feature dim, spatial dims ...]
 ConvDimensionNumbers = collections.namedtuple(
     "ConvDimensionNumbers", ["lhs_spec", "rhs_spec", "out_spec"])
 
diff --git a/tests/lax_test.py b/tests/lax_test.py
index 8c538d8f5..7da257317 100644
--- a/tests/lax_test.py
+++ b/tests/lax_test.py
@@ -422,10 +422,11 @@ class LaxTest(jtu.JaxTestCase):
       for lhs_dilation, rhs_dilation in itertools.product(
           [(1, 1), (1, 2)], repeat=2)
       for rng in [jtu.rand_small()]
-      for dim_nums, perms in [(("NCHW", "OIHW", "NCHW"),
-                              ([0, 1, 2, 3], [0, 1, 2, 3])),
-                             (("NHWC", "HWIO", "NHWC"),
-                              ([0, 2, 3, 1], [2, 3, 1, 0]))]))
+      for dim_nums, perms in [
+        (("NCHW", "OIHW", "NCHW"), ([0, 1, 2, 3], [0, 1, 2, 3])),
+        (("NHWC", "HWIO", "NHWC"), ([0, 2, 3, 1], [2, 3, 1, 0])),
+        (("NCHW", "HWIO", "NHWC"), ([0, 1, 2, 3], [2, 3, 1, 0])),
+      ]))
   def testConvGeneralDilated(self, lhs_shape, rhs_shape, dtype, strides,
                              padding, lhs_dilation, rhs_dilation,
                              dimension_numbers, perms, rng):
@@ -1756,7 +1757,8 @@ class LaxAutodiffTest(jtu.JaxTestCase):
       for rng in [jtu.rand_default()]
       for dim_nums, perms in [
           (("NCHW", "OIHW", "NCHW"), ([0, 1, 2, 3], [0, 1, 2, 3])),
-          (("NHWC", "HWIO", "NHWC"), ([0, 2, 3, 1], [2, 3, 1, 0]))
+          (("NHWC", "HWIO", "NHWC"), ([0, 2, 3, 1], [2, 3, 1, 0])),
+          (("NHWC", "OIHW", "NCHW"), ([0, 2, 3, 1], [0, 1, 2, 3]))
       ]))
   @jtu.skip_on_devices("tpu")
   def testConvGeneralDilatedGrad(self, lhs_shape, rhs_shape, dtype, strides,