Repository: google/jax

Issue:
jax.eval_shape does not understand arbitrary object dtype
```python
import objax
import jax
m = objax.nn.Sequential([objax.nn.Linear(3, 4)])

v = objax.random.uniform(((32, 3)))
jax.eval_shape(m, v)  # ShapeDtypeStruct(shape=(32, 4), dtype=float32)

d = objax.util.EasyDict(shape=(32, 3), dtype=v.dtype)
print(repr(d.shape), repr(d.dtype)) # (32, 3) dtype('float32')
print(hasattr(d, 'shape'), hasattr(d, 'dtype')) # True True

jax.eval_shape(m, d)
# 
Traceback (most recent call last):
  File "/usr/local/google/home/dberth/jax3/lib/python3.8/site-packages/IPython/core/interactiveshell.py", line 3417, in run_code
    exec(code_obj, self.user_global_ns, self.user_ns)
  File "<ipython-input-24-d48e201250f2>", line 1, in <module>
    jax.eval_shape(m, d)
  File "/usr/local/google/home/dberth/jax3/lib/python3.8/site-packages/jax/api.py", line 2283, in eval_shape
    *map(abstractify, args_flat))
  File "/usr/local/google/home/dberth/jax3/lib/python3.8/site-packages/jax/util.py", line 36, in safe_map
    return list(map(f, *args))
  File "/usr/local/google/home/dberth/jax3/lib/python3.8/site-packages/jax/api.py", line 2279, in abstractify
    return ShapedArray(np.shape(x), dtypes.result_type(x))
  File "/usr/local/google/home/dberth/jax3/lib/python3.8/site-packages/jax/dtypes.py", line 294, in result_type
    return canonicalize_dtype(dtype(args[0]))
  File "/usr/local/google/home/dberth/jax3/lib/python3.8/site-packages/jax/dtypes.py", line 288, in dtype
    return np.result_type(x)
  File "<__array_function__ internals>", line 5, in result_type
  File "/usr/local/google/home/dberth/jax3/lib/python3.8/site-packages/numpy/core/_internal.py", line 64, in _usefields
    names, formats, offsets, titles = _makenames_list(adict, align)
  File "/usr/local/google/home/dberth/jax3/lib/python3.8/site-packages/numpy/core/_internal.py", line 40, in _makenames_list
    format = dtype(obj[0], align=align)
TypeError: data type not understood
```

The class EasyDict is simply this:
```python
class EasyDict(dict):
    """Custom dictionary that allows to access dict values as attributes."""

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.__dict__ = self
```




Discussion:
Thanks for surfacing this!

JAX in general does not support the object dtype. I think we could make the error message clearer there. But actually I think the main issue is something else.

I think the main issue here is that `eval_shape` isn't living up to its contract of just requiring that its input has `shape` and `dtype` attributes. That is a bug, either in how we're calling `np.shape` or how we've implemented `dtypes.result_type` (both called in `jax.eval_shape`).

By the way, is the main goal with `EasyDict` to set up something like [`types.SimpleNamespace`](https://docs.python.org/3/library/types.html#types.SimpleNamespace)? I believe `eval_shape` works with those.

Resolving patch:
diff --git a/jax/api.py b/jax/api.py
--- a/jax/api.py
+++ b/jax/api.py
@@ -2318,8 +2318,14 @@ def __init__(self, shape, dtype):
   >>> print(out.dtype)
   float32
   """
+  def dtype(x):
+    try:
+      return dtypes.result_type(x)
+    except ValueError:
+      return dtypes.result_type(getattr(x, 'dtype'))
+
   def abstractify(x):
-    return ShapedArray(np.shape(x), dtypes.result_type(x))
+    return ShapedArray(np.shape(x), dtype(x))
   args_flat, in_tree = tree_flatten((args, kwargs))
   wrapped_fun, out_tree = flatten_fun(lu.wrap_init(fun), in_tree)
   out = pe.abstract_eval_fun(wrapped_fun.call_wrapped,
diff --git a/jax/dtypes.py b/jax/dtypes.py
--- a/jax/dtypes.py
+++ b/jax/dtypes.py
@@ -322,10 +322,12 @@ def dtype(x):
   return np.result_type(x)
 
 def _result_type_raw(*args):
-  if len(args) < 2:
+  if len(args) == 1:
     return _jax_type(args[0])
   return _least_upper_bound(*{_jax_type(arg) for arg in args})
 
 def result_type(*args):
   """Convenience function to apply Numpy argument dtype promotion."""
+  if len(args) == 0:
+    raise ValueError("at least one array or dtype is required")
   return canonicalize_dtype(_result_type_raw(*args))


Test patch:
diff --git a/tests/api_test.py b/tests/api_test.py
--- a/tests/api_test.py
+++ b/tests/api_test.py
@@ -1222,6 +1222,17 @@ def __init__(self, shape, dtype):
 
     self.assertEqual(out_shape.shape, (3, 5))
 
+  def test_eval_shape_duck_typing2(self):
+    # https://github.com/google/jax/issues/5683
+    class EasyDict(dict):
+      def __init__(self, *args, **kwargs):
+          super().__init__(*args, **kwargs)
+          self.__dict__ = self
+
+    x = EasyDict(shape=(3,), dtype=np.dtype('float32'))
+    out_shape = api.eval_shape(lambda x: x, x)  # doesn't crash
+    self.assertEqual(out_shape.shape, (3,))
+
   def test_issue_871(self):
     T = jnp.array([[1., 2.], [3., 4.], [5., 6.]])
     x = jnp.array([1, 2, 3])