Repository: google/jax

Issue:
jnp.array breaks JIT invariance for some list inputs
Minimal repro:
```python
import numpy as np
import jax.numpy as jnp
from jax import jit

x = [0, np.int16(1)]

print(jnp.array(x).dtype)
# int32

print(jit(jnp.array)(x).dtype)
# int16
```
The reason for this is that, outside JIT, `jnp.array(x)` falls back to `np.array(x)`, presumably for performance reasons, which results in numpy-style type promotion being applied to the list contents in place of JAX-style type promotion: https://github.com/google/jax/blob/2f4bd118839b29bdf9aa84f8c5b351ce6d7bff71/jax/_src/numpy/lax_numpy.py#L3589-L3594

Within JIT, however, the list contents are traced and `_can_call_numpy_array` evaluates to False, meaning that the final dtype is determined by `jnp.stack`, which uses JAX type promotion to decide on the output type: https://github.com/google/jax/blob/2f4bd118839b29bdf9aa84f8c5b351ce6d7bff71/jax/_src/numpy/lax_numpy.py#L3614-L3616

I think the best fix here would be to use JAX-style type promotion in both cases.


Discussion:

Resolving patch:
diff --git a/jax/_src/numpy/lax_numpy.py b/jax/_src/numpy/lax_numpy.py
--- a/jax/_src/numpy/lax_numpy.py
+++ b/jax/_src/numpy/lax_numpy.py
@@ -475,25 +475,6 @@ def _jnp_dtype(obj, align=False, copy=False):
   64: np.int64,
 }
 
-def _np_array(obj, dtype=None, **kwargs):
-  """Return a properly-typed numpy array.
-
-  `_np_array(obj, **kwds)` is equivalent to `np.array(obj, **kwds)`, with the
-  exception that when obj.dtype is not defined and dtype is not specified, it
-  uses Jax's default dtypes.
-  """
-  arr = np.array(obj, dtype=dtype, **kwargs)
-  obj_dtype = getattr(obj, 'dtype', None)
-  arr_dtype = np.dtype(arr.dtype).type
-  if dtype is None and obj_dtype is None:
-    if dtypes.is_python_scalar(obj):
-      arr = arr.astype(result_type(obj))
-    elif arr_dtype in _DEFAULT_TYPEMAP:
-      arr = arr.astype(_DEFAULT_TYPEMAP[arr_dtype])
-  return arr
-
-_np_asarray = partial(_np_array, copy=False)
-
 def _promote_shapes(fun_name, *args):
   """Prepend implicit leading singleton dimensions for Numpy broadcasting."""
   if len(args) < 2:
@@ -3574,7 +3555,6 @@ def atleast_3d(*arys):
 https://jax.readthedocs.io/en/latest/faq.html).
 """
 
-
 @_wraps(np.array, lax_description=_ARRAY_DOC)
 def array(object, dtype=None, copy=True, order="K", ndmin=0):
   if order is not None and order != "K":
@@ -3583,28 +3563,45 @@ def array(object, dtype=None, copy=True, order="K", ndmin=0):
   # check if the given dtype is compatible with JAX
   lax._check_user_dtype_supported(dtype, "array")
 
+  # Here we make a judgment call: we only return a weakly-typed array when the
+  # input object itself is weakly typed. That ensures asarray(x) is a no-op whenever
+  # x is weak, but avoids introducing weak types with something like array([1, 2, 3])
   weak_type = dtype is None and dtypes.is_weakly_typed(object)
-  dtype = dtype and dtypes.canonicalize_dtype(dtype)
 
-  if _can_call_numpy_array(object):
-    if dtypes.is_python_scalar(object):
-      object = dtypes.coerce_to_array(object, dtype)
-    # TODO(jakevdp): falling back to numpy here fails to overflow for lists containing
-    # large integers; see discussion in https://github.com/google/jax/pull/6047.
-    object = _np_array(object, dtype=dtype, ndmin=ndmin, copy=False)
+  # For Python scalar literals, call coerce_to_array to catch any overflow errors.
+  # We don't use dtypes.is_python_scalar because we don't want this triggering for
+  # traced values. We do this here because it matters whether or not dtype is None.
+  # We don't assign the result because we want the raw object to be used for type
+  # inference below.
+  if isinstance(object, (bool, int, float, complex)):
+    _ = dtypes.coerce_to_array(object, dtype)
+
+  leaves = tree_leaves(object)
+  if dtype is None:
+    # Use lattice_result_type rather than result_type to avoid canonicalization.
+    # Otherwise, weakly-typed inputs would have their dtypes canonicalized.
+    try:
+      dtype = dtypes._lattice_result_type(*leaves)[0] if leaves else dtypes.float_
+    except TypeError:
+      # This happens if, e.g. one of the entries is a memoryview object.
+      # This is rare, so we only handle it if the normal path fails.
+      leaves = [_convert_to_array_if_dtype_fails(leaf) for leaf in leaves]
+      dtype = dtypes._lattice_result_type(*leaves)[0]
 
-    # call _np_array a second time with canonicalized dtype
-    dtype = dtypes.canonicalize_dtype(object.dtype)
-    object = _np_array(object, dtype=dtype, copy=False)
+  if not weak_type:
+    dtype = dtypes.canonicalize_dtype(dtype)
 
-  assert type(object) not in dtypes.python_scalar_dtypes
+  # We can't use the ndarray class because we need to handle internal buffers
+  # (See https://github.com/google/jax/issues/8950)
+  ndarray_types = (device_array.DeviceArray, core.Tracer)
 
-  if type(object) is np.ndarray:
-    _inferred_dtype = object.dtype and dtypes.canonicalize_dtype(object.dtype)
-    lax._check_user_dtype_supported(_inferred_dtype, "array")
-    out = _np_array(object, copy=copy, dtype=dtype)
-    if dtype: assert _dtype(out) == dtype
-  elif isinstance(object, (device_array.DeviceArray, core.Tracer)):
+  if not _any(isinstance(leaf, ndarray_types) for leaf in leaves):
+    # TODO(jakevdp): falling back to numpy here fails to overflow for lists containing
+    # large integers; see discussion in https://github.com/google/jax/pull/6047.
+    # More correct would be to call coerce_to_array on each leaf, but this may have
+    # performance implications.
+    out = np.array(object, dtype=dtype, ndmin=ndmin, copy=False)
+  elif isinstance(object, ndarray_types):
     if object.aval is None:
       # object is a raw buffer; convert to device array on its current device.
       aval = ShapedArray(object.xla_shape().dimensions(), object.dtype,
@@ -3615,34 +3612,30 @@ def array(object, dtype=None, copy=True, order="K", ndmin=0):
     if object:
       out = stack([asarray(elt, dtype=dtype) for elt in object])
     else:
-      out = _np_array([], dtype=dtype)
+      out = np.array([], dtype=dtype)
   else:
     try:
       view = memoryview(object)
     except TypeError:
       pass  # `object` does not support the buffer interface.
     else:
-      return array(_np_asarray(view), dtype, copy, ndmin=ndmin)
+      return array(np.asarray(view), dtype, copy, ndmin=ndmin)
 
     raise TypeError("Unexpected input type for array: {}".format(type(object)))
 
-  if weak_type:
-    # Here we make a judgment call: we only return a weakly-typed array when obj
-    # itself is weakly typed. That ensures array(x) is a no-op whenever x is weak,
-    # but avoids introducing weak types with something like array([1, 2, 3])
-    out = lax._convert_element_type(out, dtype, weak_type=True)
-  else:
-    # If dtype is not specified, we use result_type(out). This ensures JIT invariance
-    # with, e.g. lists of scalars.
-    out = lax._convert_element_type(out, dtype or result_type(out))
-
+  out = lax._convert_element_type(out, dtype, weak_type=weak_type)
   if ndmin > ndim(out):
-    out = lax.broadcast(out, (1,) * (ndmin - ndim(out)))
+    out = lax.expand_dims(out, range(ndmin - ndim(out)))
   return out
 
-def _can_call_numpy_array(x):
-  return _all(not isinstance(l, (core.Tracer, device_array.DeviceArray))
-              for l in tree_leaves(x))
+
+def _convert_to_array_if_dtype_fails(x):
+  try:
+    dtypes.dtype(x)
+  except TypeError:
+    return np.asarray(x)
+  else:
+    return x
 
 
 @_wraps(np.asarray, lax_description=_ARRAY_DOC)


Test patch:
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
@@ -3667,6 +3667,11 @@ def _check(obj, out_dtype, weak_type):
     _check([jnp.float64(1)], np.float64, False)
     _check([jnp.complex128(1)], np.complex128, False)
 
+    # Mixed inputs use JAX-style promotion.
+    # (regression test for https://github.com/google/jax/issues/8945)
+    _check([0, np.int16(1)], np.int16, False)
+    _check([0.0, np.float16(1)], np.float16, False)
+
   @parameterized.named_parameters(jtu.cases_from_list(
       {"testcase_name": f"_dtype={np.dtype(dtype)}", "dtype": dtype}
       for dtype in all_dtypes))
@@ -3753,12 +3758,12 @@ def __array__(self, dtype=None):
 
   def testArrayMethod(self):
     class arraylike(object):
-      dtype = np.float32
+      dtype = np.dtype('float32')
       def __array__(self, dtype=None):
         return np.array(3., dtype=dtype)
     a = arraylike()
     ans = jnp.array(a)
-    assert ans == 3.
+    self.assertEqual(ans, 3.)
 
   def testMemoryView(self):
     self.assertAllClose(