add test

jax/abstract_arrays.py

@@ -68,9 +68,8 @@ def _zeros_like_python_scalar(t, x):
  return np.array(0, dtypes.python_scalar_dtypes[t])
 
 def _make_concrete_python_scalar(t, x):
- return ConcreteArray(
- np.array(x, dtype=dtypes.python_scalar_dtypes[t]),
- weak_type=True)
+ return ConcreteArray(np.array(x, dtype=dtypes.python_scalar_dtypes[t]),
+ weak_type=True)
 
 for t in dtypes.python_scalar_dtypes:
  core.pytype_aval_mappings[t] = partial(_make_concrete_python_scalar, t)

tests/api_test.py

@@ -2393,6 +2393,18 @@ def f(_):
  expected = jnp.arange(1) + 1
  self.assertAllClose(ans, expected)
 
+ def test_large_python_int_to_float(self):
+ # https://github.com/google/jax/pull/6165
+ # We skip checks because otherwise we end up calling valid_jaxtype(2**100),
+ # which tries to form a ConcreteArray with that value and thus leads to a
+ # NumPy OverflowError. It's true that 2**100 does not inhabit a jax type,
+ # but as an issue of Python embedding we can handle operations like those in
+ # the tests below.
+ with jax.core.skipping_checks():
+ jnp.multiply(2 ** 100, 3.) # doesn't crash
+ out = lax.convert_element_type(2 ** 100, jnp.float32) # doesn't crash
+ self.assertArraysEqual(out, np.float32(2 ** 100))
+
 
 class RematTest(jtu.JaxTestCase):