Repository: google/jax

Issue:
jax.ShapeDtypeStruct isn't documented
It is helpful for AOT compilation!


Discussion:
We should also probably update the documentation of [`jax.eval_shape`](https://jax.readthedocs.io/en/latest/_autosummary/jax.eval_shape.html) to make use of `jax.ShapeDtypeStruct` rather than suggesting that users define a similar structure themselves.
We use `types.SimpleNamespace` to a similar end as an example throughout our reference documentation, e.g. [here](https://github.com/google/jax/blob/d5f8fd3d851acca570cc588ee4a8ef833a9fe7c1/jax/_src/api.py#L2338).

On the one hand, this isn't `ShapeDtypeStruct`. On the other hand, it helps emphasize that anything that duck-types is meant to work. Should we keep those examples as is, or give up on emphasizing the duck-typing point as part of the docstring example, and switch it over to `ShapeDtypeStruct` as well?
We should either document `ShapeDtypeStruct` or remove it from the API.

If we expect users to use duck types, we should describe clearly what attributes the duck types are expected to have.
Do you mean to state those as mutually exclusive? It seems at least possible to have SDS, document it, allow for anything that duck types like SDS on the input side, and always output SDS. With that, it might make sense to have the SDS symbol in the API for type annotations. (?)

We could entertain a protocol for annotations, but by the time we've written a class to define the protocol, it seems easy to fill it out into a concrete class for optional use too?
I didn't mean they were exclusive. I just mean without documentation, it's hard to work out what attributes matter, and maybe the best form of that documentation is a class like SDS.
I agree with @hawkinsp – using `ShapeDtypeStruct` in the docs is helpful in that it's self-documenenting. We could also mention that any duck-typed object will work. but defining a struct in the example makes it look more complicated than it needs to be.
Yeah, great, I agree!

Resolving patch:
diff --git a/jax/_src/api.py b/jax/_src/api.py
--- a/jax/_src/api.py
+++ b/jax/_src/api.py
@@ -2704,6 +2704,16 @@ def device_get(x: Any):
 
 
 class ShapeDtypeStruct:
+  """A container for the shape, dtype, and other static attributes of an array.
+
+  ``ShapeDtypeStruct`` is often used in conjunction with :func:`jax.eval_shape`.
+
+  Args:
+    shape: a sequence of integers representing an array shape
+    dtype: a dtype-like object
+    named_shape: (optional) a dictionary representing a named shape
+    sharding: (optional) a :class:`jax.Sharding` object
+  """
   __slots__ = ["shape", "dtype", "named_shape", "sharding"]
   def __init__(self, shape, dtype, named_shape=None, sharding=None):
     self.shape = tuple(shape)
@@ -2764,20 +2774,9 @@ def eval_shape(fun: Callable, *args, **kwargs):
 
     def eval_shape(fun, *args, **kwargs):
       out = fun(*args, **kwargs)
+      shape_dtype_struct = lambda x: jax.ShapeDtypeStruct(x.shape, x.dtype)
       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.
@@ -2790,26 +2789,24 @@ def __init__(self, shape, dtype):
     *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.
+      accessed, one can use :class:`jax.ShapeDtypeStruct` or another container
+      that duck-types as ndarrays (note however that duck-typed objects cannot
+      be namedtuples because those are treated as standard Python containers).
     **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.
 
+  Returns:
+    out: a nested PyTree containing :class:`jax.ShapeDtypeStruct` objects as leaves.
+
   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 = jnp.dtype(dtype)
-  ...
-  >>> A = MyArgArray((2000, 3000), jnp.float32)
-  >>> x = MyArgArray((3000, 1000), jnp.float32)
+  >>> A = jax.ShapeDtypeStruct((2000, 3000), jnp.float32)
+  >>> x = jax.ShapeDtypeStruct((3000, 1000), jnp.float32)
   >>> out = jax.eval_shape(f, A, x)  # no FLOPs performed
   >>> print(out.shape)
   (2000, 1000)


Test patch: