Repository: google/jax

Issue:
Open source tests
We'd like contributors to be able to make sure their code isn't breaking anything, and that they add tests that cover their own contributions.


Discussion:

Resolving patch:
diff --git a/jax/lax.py b/jax/lax.py
--- a/jax/lax.py
+++ b/jax/lax.py
@@ -1027,8 +1027,8 @@ def conv_general_dilated_transpose_rhs(
     dimension_numbers, lhs_shape, rhs_shape):
   assert type(dimension_numbers) is ConvDimensionNumbers
   lhs_sdims, rhs_sdims, out_sdims = map(_conv_sdims, dimension_numbers)
-  transposed = map(_conv_transpose, dimension_numbers)
-  trans_dimension_numbers = ConvDimensionNumbers(*transposed)
+  lhs_trans, rhs_trans, out_trans = map(_conv_transpose, dimension_numbers)
+  trans_dimension_numbers = ConvDimensionNumbers(lhs_trans, out_trans, rhs_trans)
   padding = _conv_general_vjp_rhs_padding(
       onp.take(lhs_shape, lhs_sdims), onp.take(rhs_shape, rhs_sdims),
       window_strides, onp.take(g.shape, out_sdims), padding, lhs_dilation,


Test patch:
diff --git a/tests/examples_test.py b/examples/examples_test.py
similarity index 100%
rename from tests/examples_test.py
rename to examples/examples_test.py
diff --git a/jax/test_util.py b/jax/test_util.py
--- a/jax/test_util.py
+++ b/jax/test_util.py
@@ -46,7 +46,7 @@
 
 flags.DEFINE_integer(
   'num_generated_cases',
-  os.getenv('JAX_NUM_GENERATED_CASES', 100),
+  os.getenv('JAX_NUM_GENERATED_CASES', 10),
   help='Number of generated cases to test')
 
 EPS = 1e-4
@@ -71,6 +71,7 @@ def numpy_close(a, b, atol=ATOL, rtol=RTOL, equal_nan=False):
   if testing_tpu or testing_x32:
     atol = max(atol, 1e-1)
     rtol = max(rtol, 1e-1)
+  assert a.shape == b.shape
   return onp.allclose(a, b, atol=atol * a.size, rtol=rtol * b.size,
                       equal_nan=equal_nan)
 
diff --git a/tests/lax_numpy_test.py b/tests/lax_numpy_test.py
--- a/tests/lax_numpy_test.py
+++ b/tests/lax_numpy_test.py
@@ -139,7 +139,7 @@ def op_record(name, nargs, dtypes, shapes, rng, diff_modes, test_name=None):
     op_record("prod", 1, default_dtypes, all_shapes, jtu.rand_small_positive(), []),
     op_record("sum", 1, default_dtypes, all_shapes, jtu.rand_default(), []),
     op_record("var", 1, default_dtypes, nonempty_shapes, jtu.rand_default(), []),
-    op_record("std", 1, default_dtypes, nonempty_shapes, jtu.rand_default(), []),
+    op_record("std", 1, float_dtypes, nonempty_shapes, jtu.rand_default(), []),
 ]
 
 JAX_REDUCER_NO_DTYPE_RECORDS = [
diff --git a/tests/lax_test.py b/tests/lax_test.py
--- a/tests/lax_test.py
+++ b/tests/lax_test.py
@@ -1580,10 +1580,13 @@ def testConvWithGeneralPaddingGrad(self, lhs_shape, rhs_shape, dtype, strides,
        "rhs_dil": rhs_dil, "rng": rng, "dimension_numbers": dim_nums,
        "perms": perms}
       for lhs_shape, rhs_shape, all_strides, all_pads, lhs_dils, rhs_dils in [
-          ((b, i, 5, 6), (j, i, 1, 2), [(1, 1), (1, 2), (2, 1)],
-           [((0, 0), (0, 0)), ((1, 0), (0, 1)), ((0, -1), (0, 0))],
-           [(1, 1), (2, 1)], [(1, 1)])
-          for b, i, j in itertools.product([2, 3], repeat=3)]
+          ((b, i, 6, 7),  # lhs_shape
+           (j, i, 1, 2),  # rhs_shape
+           [(1, 1), (1, 2), (2, 1)],  # strides
+           [((0, 0), (0, 0)), ((1, 0), (0, 1)), ((0, -1), (0, 0))],  # pads
+           [(1, 1), (2, 1)],  # lhs_dils
+           [(1, 1)])  # rhs_dils
+          for b, i, j in itertools.product([1, 2], repeat=3)]
       for strides in all_strides
       for rhs_dil in rhs_dils
       for lhs_dil in lhs_dils
@@ -1592,7 +1595,7 @@ def testConvWithGeneralPaddingGrad(self, lhs_shape, rhs_shape, dtype, strides,
       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]))
       ]))
   @jtu.skip_on_devices("tpu")
   def testConvGeneralDilatedGrad(self, lhs_shape, rhs_shape, dtype, strides,