jax.lax: raise TypeError for mismatched dtypes

Author: jakevdp

jax/_src/lax/lax.py

@@ -2209,7 +2209,7 @@ def broadcasting_sharding_rule(name, *avals):
 
 
 def naryop(result_dtype, accepted_dtypes, name, allow_extended_dtype=False,
-           require_same_dtypes=False):
+           require_same_dtypes=True):
   dtype_rule = partial(naryop_dtype_rule, result_dtype, accepted_dtypes, name,
                        allow_extended_dtype=allow_extended_dtype,
                        require_same=require_same_dtypes)

tests/lax_test.py

@@ -3366,6 +3366,12 @@ def test_ops_do_not_accept_complex_dtypes(self, op):
     with self.assertRaisesRegex(TypeError, ".*does not accept dtype complex.*"):
       op(2+3j, 4+5j)
 
+  @parameterized.parameters([lax.add, lax.mul, lax.div, lax.rem, lax.lt, lax.gt,
+                             lax.ge, lax.le, lax.eq, lax.ne])
+  def test_ops_error_on_mismatched_dtypes(self, op):
+    with self.assertRaisesRegex(TypeError, ".*requires arguments to have the same dtypes.*"):
+      op(0, 0.0)
+
   def test_population_count_booleans_not_supported(self):
     # https://github.com/jax-ml/jax/issues/3886
     msg = "population_count does not accept dtype bool"