Copybara import of the project:

-- d42fffd by Jake VanderPlas <jakevdp@google.com>: JaxTestCase: set numpy_rank_promotion='raise' by default PiperOrigin-RevId: 427896974
google · Feb 11, 2022 · 5691010 · 5691010
1 parent 8b4a7ce
commit 5691010
Show file tree

Hide file tree

Showing 29 changed files with 73 additions and 20 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -29,14 +29,6 @@ PLEASE REMEMBER TO CHANGE THE '..main' WITH AN ACTUAL TAG in GITHUB LINK.
     `dialect=` is passed.
   * The `jax.jit(f).lower(...).compiler_ir(dialect='mhlo')` now returns an MLIR
     `ir.Module` object instead of its string representation.
-  * `jax.test_util.JaxTestCase` now sets `jax_numpy_rank_promotion='raise'` by
-    default. To recover the previous behavior, use the `jax.test_util.with_config`
-    decorator:
-    ```python
-    @jtu.with_config(jax_numpy_rank_promotion='allow')
-    class MyTest(jtu.JaxTestCase):
-      ...
-    ```
 
 ## jaxlib 0.1.76 (Jan 27, 2022)
 

diff --git a/jax/_src/test_util.py b/jax/_src/test_util.py
@@ -915,10 +915,7 @@ def decorator(cls):
 
 class JaxTestCase(parameterized.TestCase):
   """Base class for JAX tests including numerical checks and boilerplate."""
-  _default_config = {
-    'jax_enable_checks': True,
-    'jax_numpy_rank_promotion': 'raise',
-  }
+  _default_config = {'jax_enable_checks': True}
 
   # TODO(mattjj): this obscures the error messages from failures, figure out how
   # to re-enable it

diff --git a/jax/experimental/jax2tf/tests/tf_test_util.py b/jax/experimental/jax2tf/tests/tf_test_util.py
@@ -150,7 +150,6 @@ def ComputeTfValueAndGrad(tf_f: Callable, tf_args: Sequence,
   return f1(*args1)
 
 
-@jtu.with_config(jax_numpy_rank_promotion="allow")
 class JaxToTfTestCase(jtu.JaxTestCase):
 
   def setUp(self):

diff --git a/tests/api_test.py b/tests/api_test.py
@@ -70,6 +70,7 @@
 numpy_version = tuple(map(int, np.__version__.split('.')[:3]))
 
 
+@jtu.with_config(jax_numpy_rank_promotion="raise")
 class CPPJitTest(jtu.BufferDonationTestCase):
   """Shared tests between the Python and the C++ jax,jit implementations.
 
@@ -859,14 +860,15 @@ def f(k):
       python_should_be_executing = False
       self.assertEqual(x, f(x))
 
-
+@jtu.with_config(jax_numpy_rank_promotion="raise")
 class PythonJitTest(CPPJitTest):
 
   @property
   def jit(self):
     return api._python_jit
 
 
+@jtu.with_config(jax_numpy_rank_promotion="raise")
 class APITest(jtu.JaxTestCase):
 
   def test_grad_item(self):
@@ -3414,6 +3416,7 @@ def f(x):
         FLAGS.jax_numpy_rank_promotion = allow_promotion
 
 
+@jtu.with_config(jax_numpy_rank_promotion="raise")
 class RematTest(jtu.JaxTestCase):
 
   @parameterized.named_parameters(
@@ -4270,6 +4273,7 @@ def f(u, x):
 
     _ = api.linearize(partial(f, core.unit), 3.)
 
+@jtu.with_config(jax_numpy_rank_promotion="raise")
 class JaxprTest(jtu.JaxTestCase):
 
   def test_scalar_literals(self):
@@ -4413,6 +4417,7 @@ def test_convert_element_type_literal_constant_folding(self):
     self.assertLen(jaxpr.eqns, 0)
 
 
+@jtu.with_config(jax_numpy_rank_promotion="raise")
 class CustomJVPTest(jtu.JaxTestCase):
 
   def test_basic(self):
@@ -5387,6 +5392,7 @@ def f_jvp(primals, tangents):
     self.assertEqual(shape, ())
 
 
+@jtu.with_config(jax_numpy_rank_promotion="raise")
 class CustomVJPTest(jtu.JaxTestCase):
 
   def test_basic(self):
@@ -6355,6 +6361,7 @@ def transposed(y):
   return transposed
 
 
+@jtu.with_config(jax_numpy_rank_promotion="raise")
 class CustomTransposeTest(jtu.JaxTestCase):
 
   def test_linear_call(self):
@@ -6683,6 +6690,7 @@ def tp(r, t): return t / r
     self.assertAllClose(f_t(x), jax.jit(f_t)(x))
 
 
+@jtu.with_config(jax_numpy_rank_promotion="raise")
 class CustomVmapTest(jtu.JaxTestCase):
 
   def test_basic(self):
@@ -7109,6 +7117,7 @@ def vmap_ref(xs, y):
     self.assertEqual(str(jaxpr), str(jaxpr_ref))
 
 
+@jtu.with_config(jax_numpy_rank_promotion="raise")
 class CustomApiTest(jtu.JaxTestCase):
   """Test interactions among the custom_{vmap,jvp,vjp,transpose,*} APIs"""
 
@@ -7146,6 +7155,7 @@ def test_def_method_forwarding_all_permutations(self):
           self.assertIsInstance(getattr(f, method), Callable)
 
 
+@jtu.with_config(jax_numpy_rank_promotion="raise")
 class InvertibleADTest(jtu.JaxTestCase):
 
   @jtu.ignore_warning(message="Values that an @invertible function closes")
@@ -7254,6 +7264,7 @@ def f(x, y):
                         check_dtypes=True)
 
 
+@jtu.with_config(jax_numpy_rank_promotion="raise")
 class BufferDonationTest(jtu.BufferDonationTestCase):
 
   @jtu.skip_on_devices("cpu")  # In/out aliasing not supported on CPU.
@@ -7276,6 +7287,7 @@ def test_pmap_nested_donate_ignored(self):
     pmap_fun(a)  # doesn't crash
 
 
+@jtu.with_config(jax_numpy_rank_promotion="raise")
 class NamedCallTest(jtu.JaxTestCase):
 
   def test_default_name(self):
@@ -7356,6 +7368,7 @@ def test_integer_overflow(self, jit_type, func):
     self.assertRaises(OverflowError, f, int_min - 1)
 
 
+@jtu.with_config(jax_numpy_rank_promotion="raise")
 class BackendsTest(jtu.JaxTestCase):
 
   @unittest.skipIf(not sys.executable, "test requires sys.executable")
@@ -7378,6 +7391,7 @@ def test_cpu_warning_suppression(self):
     assert "No GPU/TPU found" not in result.stderr.decode()
 
 
+@jtu.with_config(jax_numpy_rank_promotion="raise")
 class CleanupTest(jtu.JaxTestCase):
   def test_call_wrapped_second_phase_cleanup(self):
     try:
@@ -7538,7 +7552,6 @@ def f(x, n):
     self.assertIs(jaxpr.jaxpr.invars[1], jaxpr.out_avals[0].shape[0])
     self.assertEqual(4, jaxpr.out_avals[0].shape[1])
 
-  @jax.numpy_rank_promotion("allow")  # explicitly exercises implicit rank promotion.
   def test_basic_batchpoly_neuralnet(self):
     def predict(params, inputs):
       for W, b in params:

diff --git a/tests/batching_test.py b/tests/batching_test.py
@@ -40,6 +40,7 @@
 # These are 'manual' tests for batching (vmap). The more exhaustive, more
 # systematic tests are in lax_test.py's LaxVmapTest class.
 
+@jtu.with_config(jax_numpy_rank_promotion="raise")
 class BatchingTest(jtu.JaxTestCase):
 
   def testConstantFunction(self):

diff --git a/tests/fft_test.py b/tests/fft_test.py
@@ -93,6 +93,7 @@ def _zero_for_irfft(z, axes):
   return jnp.concatenate(parts, axis=axis)
 
 
+@jtu.with_config(jax_numpy_rank_promotion="raise")
 class FftTest(jtu.JaxTestCase):
 
   def testNotImplemented(self):

diff --git a/tests/infeed_test.py b/tests/infeed_test.py
@@ -26,6 +26,7 @@
 
 config.parse_flags_with_absl()
 
+@jtu.with_config(jax_numpy_rank_promotion="raise")
 class InfeedTest(jtu.JaxTestCase):
 
   @jax.numpy_rank_promotion("allow")  # Test explicitly exercises implicit rank promotion.

diff --git a/tests/lax_autodiff_test.py b/tests/lax_autodiff_test.py
@@ -189,7 +189,7 @@ def check_grads_bilinear(f, args, order,
   check_grads(lambda rhs: f(lhs, rhs), (rhs,), order,
               modes=modes, atol=atol, rtol=rtol, eps=1.)
 
-
+@jtu.with_config(jax_numpy_rank_promotion="raise")
 class LaxAutodiffTest(jtu.JaxTestCase):
 
   @parameterized.named_parameters(itertools.chain.from_iterable(

diff --git a/tests/lax_numpy_einsum_test.py b/tests/lax_numpy_einsum_test.py
@@ -30,6 +30,7 @@
 config.parse_flags_with_absl()
 
 
+@jtu.with_config(jax_numpy_rank_promotion="raise")
 class EinsumTest(jtu.JaxTestCase):
 
   def _check(self, s, *ops):

diff --git a/tests/lax_numpy_indexing_test.py b/tests/lax_numpy_indexing_test.py
@@ -414,6 +414,7 @@ def check_grads(f, args, order, atol=None, rtol=None, eps=None):
 MODES = ["clip", "drop", "promise_in_bounds"]
 
 
+@jtu.with_config(jax_numpy_rank_promotion="raise")
 class IndexingTest(jtu.JaxTestCase):
   """Tests for Numpy indexing translation rules."""
 
@@ -996,6 +997,7 @@ def _update_tol(op):
   return tol
 
 
+@jtu.with_config(jax_numpy_rank_promotion="raise")
 class IndexedUpdateTest(jtu.JaxTestCase):
 
   @parameterized.named_parameters(jtu.named_cases_from_sampler(lambda s: ({

diff --git a/tests/lax_numpy_test.py b/tests/lax_numpy_test.py
@@ -519,6 +519,7 @@ def wrapper(*args, **kw):
   return wrapper
 
 
+@jtu.with_config(jax_numpy_rank_promotion="raise")
 class LaxBackedNumpyTests(jtu.JaxTestCase):
   """Tests for LAX-backed Numpy implementation."""
 
@@ -5914,6 +5915,7 @@ def grad_test_spec(op, nargs, order, rng_factory, dtypes, name=None, tol=None):
     GradSpecialValuesTestSpec(jnp.sinc, [0.], 1),
 ]
 
+@jtu.with_config(jax_numpy_rank_promotion="raise")
 class NumpyGradTests(jtu.JaxTestCase):
 
   @parameterized.named_parameters(itertools.chain.from_iterable(
@@ -6018,6 +6020,7 @@ def testGradLogaddexp2Complex(self, shapes, dtype):
       tol = 3e-2
     check_grads(jnp.logaddexp2, args, 1, ["fwd", "rev"], tol, tol)
 
+@jtu.with_config(jax_numpy_rank_promotion="raise")
 class NumpySignaturesTest(jtu.JaxTestCase):
 
   def testWrappedSignaturesMatch(self):
@@ -6133,6 +6136,7 @@ def _dtypes_for_ufunc(name: str) -> Iterator[Tuple[str, ...]]:
       yield arg_dtypes
 
 
+@jtu.with_config(jax_numpy_rank_promotion="raise")
 class NumpyUfuncTests(jtu.JaxTestCase):
 
   @parameterized.named_parameters(
@@ -6164,6 +6168,7 @@ def testUfuncInputTypes(self, name, arg_dtypes):
     # that jnp returns float32. e.g. np.cos(np.uint8(0))
     self._CheckAgainstNumpy(np_op, jnp_op, args_maker, check_dtypes=False, tol=1E-2)
 
+@jtu.with_config(jax_numpy_rank_promotion="raise")
 class NumpyDocTests(jtu.JaxTestCase):
 
   def test_lax_numpy_docstrings(self):

diff --git a/tests/lax_numpy_vectorize_test.py b/tests/lax_numpy_vectorize_test.py
@@ -25,6 +25,7 @@
 config.parse_flags_with_absl()
 
 
+@jtu.with_config(jax_numpy_rank_promotion="raise")
 class VectorizeTest(jtu.JaxTestCase):
 
   @parameterized.named_parameters(jtu.cases_from_list(

diff --git a/tests/lax_scipy_sparse_test.py b/tests/lax_scipy_sparse_test.py
@@ -64,6 +64,7 @@ def rand_sym_pos_def(rng, shape, dtype):
   return matrix @ matrix.T.conj()
 
 
+@jtu.with_config(jax_numpy_rank_promotion="raise")
 class LaxBackedScipyTests(jtu.JaxTestCase):
 
   def _fetch_preconditioner(self, preconditioner, A, rng=None):

diff --git a/tests/lax_scipy_test.py b/tests/lax_scipy_test.py
@@ -144,6 +144,7 @@ def op_record(name, nargs, dtypes, rng_factory, test_grad, nondiff_argnums=(), t
 ]
 
 
+@jtu.with_config(jax_numpy_rank_promotion="raise")
 class LaxBackedScipyTests(jtu.JaxTestCase):
   """Tests for LAX-backed Scipy implementation."""
 

diff --git a/tests/lax_test.py b/tests/lax_test.py
@@ -181,6 +181,7 @@ def op_record(op, nargs, dtypes, rng_factory, tol=None):
 ]
 
 
+@jtu.with_config(jax_numpy_rank_promotion="raise")
 class LaxTest(jtu.JaxTestCase):
   """Numerical tests for LAX operations."""
 
@@ -2668,6 +2669,7 @@ def testDynamicSliceU8Index(self):
         np.array(lax.dynamic_slice(x, np.uint8([128]), (1,))), [128])
 
 
+@jtu.with_config(jax_numpy_rank_promotion="raise")
 class LazyConstantTest(jtu.JaxTestCase):
   def _Check(self, make_const, expected):
     # check casting to ndarray works
@@ -2870,6 +2872,7 @@ def testLog1pNearOne(self):
         np.log1p(np.float32(1e-5)), lax.log1p(np.complex64(1e-5)))
 
 
+@jtu.with_config(jax_numpy_rank_promotion="raise")
 class LaxNamedShapeTest(jtu.JaxTestCase):
 
   def test_abstract_eval(self):

diff --git a/tests/linalg_test.py b/tests/linalg_test.py
@@ -42,6 +42,7 @@
 complex_types = jtu.dtypes.complex
 
 
+@jtu.with_config(jax_numpy_rank_promotion='raise')
 class NumpyLinalgTest(jtu.JaxTestCase):
 
   def testNotImplemented(self):
@@ -956,6 +957,7 @@ def f(inp):
     self.assertFalse(np.any(np.isnan(cube_func(a))))
 
 
+@jtu.with_config(jax_numpy_rank_promotion='raise')
 class ScipyLinalgTest(jtu.JaxTestCase):
 
   @parameterized.named_parameters(jtu.cases_from_list(
@@ -1372,7 +1374,7 @@ def expm(x):
       jtu.check_grads(expm, (a,), modes=["fwd", "rev"], order=1, atol=tol,
                       rtol=tol)
 
-
+@jtu.with_config(jax_numpy_rank_promotion='raise')
 class LaxLinalgTest(jtu.JaxTestCase):
 
   def run_test(self, alpha, beta):

diff --git a/tests/nn_test.py b/tests/nn_test.py
@@ -35,6 +35,7 @@
 config.parse_flags_with_absl()
 
 
+@jtu.with_config(jax_numpy_rank_promotion="raise")
 class NNFunctionsTest(jtu.JaxTestCase):
   @jtu.skip_on_flag("jax_skip_slow_tests", True)
   def testSoftplusGrad(self):
@@ -229,7 +230,7 @@ def initializer_record(name, initializer, dtypes, min_dims=2, max_dims=4):
     initializer_record("delta_orthogonal", nn.initializers.delta_orthogonal, jtu.dtypes.floating, 4, 4)
 ]
 
-
+@jtu.with_config(jax_numpy_rank_promotion="raise")
 class NNInitializersTest(jtu.JaxTestCase):
   @parameterized.named_parameters(jtu.cases_from_list(
       {"testcase_name":

diff --git a/tests/pjit_test.py b/tests/pjit_test.py
@@ -74,6 +74,7 @@ def check_1d_2d_mesh(f, set_mesh):
 
 
 # TODO(skye): make the buffer donation utils part of JaxTestCase
+@jtu.with_config(jax_numpy_rank_promotion="raise")
 class PJitTest(jtu.BufferDonationTestCase):
 
   @jtu.with_mesh([('x', 1)])
@@ -634,6 +635,7 @@ def f(x, y):
     self.assertEqual(f(1, 'bye'), 5)
 
 
+@jtu.with_config(jax_numpy_rank_promotion="raise")
 class GDAPjitTest(jtu.JaxTestCase):
 
   @jtu.with_mesh([('x', 4), ('y', 2)])
@@ -951,6 +953,7 @@ def spec_regex(s):
   return str(s).replace(r"(", r"\(").replace(r")", r"\)")
 
 
+@jtu.with_config(jax_numpy_rank_promotion="raise")
 class PJitErrorTest(jtu.JaxTestCase):
   @check_1d_2d_mesh(set_mesh=True)
   def testNonDivisibleArgs(self, mesh, resources):
@@ -1178,6 +1181,7 @@ def h(x):
       f(x)
 
 
+@jtu.with_config(jax_numpy_rank_promotion="raise")
 class UtilTest(jtu.JaxTestCase):
 
   def testOpShardingRoundTrip(self):