Repository: google/jax

Issue:
jnp.load() crashes if save file contains object that is not array
Minimum repro:
```python
import jax.numpy as jnp

class TestClass:
    pass

jnp.save("object.npy", TestClass())
print(jnp.load("object.npy", allow_pickle=True))
```

Output:
```
Traceback (most recent call last):
  File ".../lib/python3.9/site-packages/jax/_src/numpy/lax_numpy.py", line 3632, in array
    dtype = dtypes._lattice_result_type(*leaves)[0] if leaves else dtypes.float_
  File ".../lib/python3.9/site-packages/jax/_src/dtypes.py", line 362, in _lattice_result_type
    dtypes, weak_types = zip(*(_dtype_and_weaktype(arg) for arg in args))
  File ".../lib/python3.9/site-packages/jax/_src/dtypes.py", line 362, in <genexpr>
    dtypes, weak_types = zip(*(_dtype_and_weaktype(arg) for arg in args))
  File ".../lib/python3.9/site-packages/jax/_src/dtypes.py", line 251, in _dtype_and_weaktype
    return dtype(value), any(value is typ for typ in _weak_types) or is_weakly_typed(value)
  File ".../lib/python3.9/site-packages/jax/_src/dtypes.py", line 357, in dtype
    raise TypeError(f"Value '{x}' with dtype {dt} is not a valid JAX array "
TypeError: Value '<__main__.TestClass object at 0x7f8528b71160>' with dtype object is not a valid JAX array type. Only arrays of numeric types are supported by JAX.

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "repro.py", line 9, in <module>
    print(jnp.load("object.npy", allow_pickle=True))
  File ".../lib/python3.9/site-packages/jax/_src/numpy/lax_numpy.py", line 652, in load
    out = asarray(out)
  File ".../lib/python3.9/site-packages/jax/_src/numpy/lax_numpy.py", line 3693, in asarray
    return array(a, dtype=dtype, copy=False, order=order)
  File ".../lib/python3.9/site-packages/jax/_src/numpy/lax_numpy.py", line 3637, in array
    dtype = dtypes._lattice_result_type(*leaves)[0]
  File ".../lib/python3.9/site-packages/jax/_src/dtypes.py", line 362, in _lattice_result_type
    dtypes, weak_types = zip(*(_dtype_and_weaktype(arg) for arg in args))
  File ".../lib/python3.9/site-packages/jax/_src/dtypes.py", line 362, in <genexpr>
    dtypes, weak_types = zip(*(_dtype_and_weaktype(arg) for arg in args))
  File ".../lib/python3.9/site-packages/jax/_src/dtypes.py", line 251, in _dtype_and_weaktype
    return dtype(value), any(value is typ for typ in _weak_types) or is_weakly_typed(value)
  File ".../lib/python3.9/site-packages/jax/_src/dtypes.py", line 357, in dtype
    raise TypeError(f"Value '{x}' with dtype {dt} is not a valid JAX array "
TypeError: Value '<__main__.TestClass object at 0x7f8528b71160>' with dtype object is not a valid JAX array type. Only arrays of numeric types are supported by JAX.
```

jax: 0.3.1
jaxlib: 0.3.0
numpy: 1.21.2
Python: 3.9.7

Maybe I am doing something wrong, but in previous jax versions (e.g. jax-0.2.25) this used to work. `jnp.load` would return an array containing a single element, the Python object, which is what the numpy version does.

Potentially this commit is what changed the behaviour? 1137aa1

Let me know if you need any more information!


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
@@ -648,7 +648,10 @@ def load(*args, **kwargs):
     # numpy does not recognize bfloat16, so arrays are serialized as void16
     if out.dtype == 'V2':
       out = out.view(bfloat16)
-    out = asarray(out)
+    try:
+      out = asarray(out)
+    except TypeError:  # Unsupported dtype
+      pass
   return out
 
 ### implementations of numpy functions in terms of lax


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
@@ -538,15 +538,19 @@ def testNotImplemented(self):
         func()
 
   @parameterized.named_parameters(jtu.cases_from_list(
-      {"testcase_name": "_{}".format(dtype), "dtype": dtype}
-      for dtype in float_dtypes))
-  def testLoad(self, dtype):
+      {"testcase_name": "_{}_allow_picke={}".format(dtype, allow_pickle),
+       "dtype": dtype, "allow_pickle": allow_pickle}
+      for dtype in float_dtypes + [object]
+      for allow_pickle in [True, False]))
+  def testLoad(self, dtype, allow_pickle):
+    if dtype == object and not allow_pickle:
+      self.skipTest("dtype=object requires allow_pickle=True")
     rng = jtu.rand_default(self.rng())
     arr = rng((10), dtype)
     with io.BytesIO() as f:
       jnp.save(f, arr)
       f.seek(0)
-      arr_out = jnp.load(f)
+      arr_out = jnp.load(f, allow_pickle=allow_pickle)
     self.assertArraysEqual(arr, arr_out)
 
   @parameterized.named_parameters(itertools.chain.from_iterable(