# Buggy code
def eval_shape(fun: Callable, *args, **kwargs):
    """Compute the shape/dtype of ``fun`` without any FLOPs.

    This utility function is useful for performing shape inference. Its
    input/output behavior is defined by::

      def eval_shape(fun, *args, **kwargs):
        out = fun(*args, **kwargs)
        return jax.tree_util.tree_map(shape_dtype_struct, out)

      def shape_dtype_struct(x):
        return ShapeDtypeStruct(x.shape, x.dtype)

      class ShapeDtypeStruct:
        __slots__ = ["shape", "dtype"]
        def __init__(self, shape, dtype):
          self.shape = shape
          self.dtype = dtype

    In particular, the output is a pytree of objects that have ``shape`` and
    ``dtype`` attributes, but nothing else about them is guaranteed by the API.

    But instead of applying ``fun`` directly, which might be expensive, it uses
    JAX's abstract interpretation machinery to evaluate the shapes without doing
    any FLOPs.

    Using :py:func:`eval_shape` can also catch shape errors, and will raise same
    shape errors as evaluating ``fun(*args, **kwargs)``.

    Args:
      fun: The function whose output shape should be evaluated.
      *args: a positional argument tuple of arrays, scalars, or (nested) standard
        Python containers (tuples, lists, dicts, namedtuples, i.e. pytrees) of
        those types. Since only the ``shape`` and ``dtype`` attributes are
        accessed, only values that duck-type arrays are required, rather than real
        ndarrays. The duck-typed objects cannot be namedtuples because those are
        treated as standard Python containers. See the example below.
      **kwargs: a keyword argument dict of arrays, scalars, or (nested) standard
        Python containers (pytrees) of those types. As in ``args``, array values
        need only be duck-typed to have ``shape`` and ``dtype`` attributes.

    For example:

    >>> import jax
    >>> import jax.numpy as jnp
    >>>
    >>> f = lambda A, x: jnp.tanh(jnp.dot(A, x))
    >>> class MyArgArray(object):
    ...   def __init__(self, shape, dtype):
    ...     self.shape = shape
    ...     self.dtype = dtype
    ...
    >>> A = MyArgArray((2000, 3000), jnp.float32)
    >>> x = MyArgArray((3000, 1000), jnp.float32)
    >>> out = jax.eval_shape(f, A, x)  # no FLOPs performed
    >>> print(out.shape)
    (2000, 1000)
    >>> print(out.dtype)
    float32
    """

    def abstractify(x):
        return ShapedArray(np.shape(x), dtypes.result_type(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, *map(abstractify, args_flat))
    out = [ShapeDtypeStruct(x.shape, x.dtype) for x in out]
    return tree_unflatten(out_tree(), out)


# Traceback (TypeError)
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

# Fixed code
def eval_shape(fun: Callable, *args, **kwargs):
    """Compute the shape/dtype of ``fun`` without any FLOPs.

    This utility function is useful for performing shape inference. Its
    input/output behavior is defined by::

      def eval_shape(fun, *args, **kwargs):
        out = fun(*args, **kwargs)
        return jax.tree_util.tree_map(shape_dtype_struct, out)

      def shape_dtype_struct(x):
        return ShapeDtypeStruct(x.shape, x.dtype)

      class ShapeDtypeStruct:
        __slots__ = ["shape", "dtype"]
        def __init__(self, shape, dtype):
          self.shape = shape
          self.dtype = dtype

    In particular, the output is a pytree of objects that have ``shape`` and
    ``dtype`` attributes, but nothing else about them is guaranteed by the API.

    But instead of applying ``fun`` directly, which might be expensive, it uses
    JAX's abstract interpretation machinery to evaluate the shapes without doing
    any FLOPs.

    Using :py:func:`eval_shape` can also catch shape errors, and will raise same
    shape errors as evaluating ``fun(*args, **kwargs)``.

    Args:
      fun: The function whose output shape should be evaluated.
      *args: a positional argument tuple of arrays, scalars, or (nested) standard
        Python containers (tuples, lists, dicts, namedtuples, i.e. pytrees) of
        those types. Since only the ``shape`` and ``dtype`` attributes are
        accessed, only values that duck-type arrays are required, rather than real
        ndarrays. The duck-typed objects cannot be namedtuples because those are
        treated as standard Python containers. See the example below.
      **kwargs: a keyword argument dict of arrays, scalars, or (nested) standard
        Python containers (pytrees) of those types. As in ``args``, array values
        need only be duck-typed to have ``shape`` and ``dtype`` attributes.

    For example:

    >>> import jax
    >>> import jax.numpy as jnp
    >>>
    >>> f = lambda A, x: jnp.tanh(jnp.dot(A, x))
    >>> class MyArgArray(object):
    ...   def __init__(self, shape, dtype):
    ...     self.shape = shape
    ...     self.dtype = dtype
    ...
    >>> A = MyArgArray((2000, 3000), jnp.float32)
    >>> x = MyArgArray((3000, 1000), jnp.float32)
    >>> out = jax.eval_shape(f, A, x)  # no FLOPs performed
    >>> print(out.shape)
    (2000, 1000)
    >>> 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), 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, *map(abstractify, args_flat))
    out = [ShapeDtypeStruct(x.shape, x.dtype) for x in out]
    return tree_unflatten(out_tree(), out)