Repository: google/jax

Issue:
Should ShapeDtypeStruct.size exist?
`ShapeDtypeStruct` currently has `dtype` and `shape`.  We could add numpy's `size` property via

```
class ShapeDtypeStruct(object):
  __slots__ = ["shape", "dtype"]
  def __init__(self, shape, dtype):
    self.shape = shape
    self.dtype = dtype

  @property
  def size(self):
    return np.prod(self.shape, dtype=int)
```

Do you want this patch?


Discussion:
+1, maybe ndim too? Also it would be nice to have `__repr__` :smile: (maybe we can make this a `NamedTuple` ?)
`NamedTuples` are bad:

```
shape, dtype = struct  # WAT
```

But yes to `__repr__`; I hit the lack of that just a moment ago. :)
Yeah, I recently learned how bad namedtuples are! What confused me was this kind of behavior:

```python
Point = namedtuple('Point', ['x', 'y'])
hash(Point((2, 3), onp.float32)) == hash(ShapeDtypeStruct((2, 3), onp.float32))
```

They're really meant to be just sugar on tuples, without type tags. That's not super relevant, but I thought I'd complain about it at every opportunity.

ShapeDtypeStruct exists because I didn't want to expose much API surface here. ShapedArray does everything we want (and more), but I didn't want to surface that as part of the API.

So we could:
1. keep ShapeDtypeStruct super minimal, no convenience methods/properties
2. add these select convenience methods/properties
3. just use ShapedArray here

I don't have strong feelings between 1 and 2, and I suspect that you both as users of this API might be best able to decide between them. I am very cautious about 3 because I don't like exposing internals, even if it means some redundant code.

WDYT?
Understood re namedtuple, while we wait for dataclasses is it worth taking a dep on [attrs](https://www.attrs.org/en/stable/) to make these pod classes more consistent and easier to read (e.g. generated slots, eq, hash, immutability etc)?

I vote for (2) and agree with not exposing the internal `ShapedArray`. I think a complete implementation would have: `shape`, `dtype`, `ndim`, `size`, `__len__` and of course `__eq__` and `__hash__`.
Rather than using attrs, I would rather require Python 3.6+ with the dataclasses backport.

Resolving patch:
diff --git a/jax/api.py b/jax/api.py
--- a/jax/api.py
+++ b/jax/api.py
@@ -1955,6 +1955,30 @@ def __init__(self, shape, dtype):
     self.shape = shape
     self.dtype = dtype
 
+  size = property(lambda self: onp.prod(self.shape))
+  ndim = property(lambda self: len(self.shape))
+
+  def __len__(self):
+    try:
+      return self.shape[0]
+    except IndexError:
+      raise TypeError("len() of unsized object")  # same as numpy error
+
+  def __repr__(self):
+    return "{}(shape={}, dtype={})".format(
+        type(self).__name__, self.shape, self.dtype.dtype.name)
+
+  __str__ = __repr__
+
+  def __eq__(self, other):
+    if not isinstance(other, ShapeDtypeStruct):
+      return False
+    else:
+      return (other.shape, other.dtype) == (self.shape, self.dtype)
+
+  def __hash__(self):
+    return hash((self.shape, self.dtype))
+
 def eval_shape(fun, *args, **kwargs):
   """Compute the shape/dtype of ``fun(*args, **kwargs)`` without any FLOPs.
 


Test patch:
diff --git a/tests/api_test.py b/tests/api_test.py
--- a/tests/api_test.py
+++ b/tests/api_test.py
@@ -891,6 +891,26 @@ def f(pt):
     g = api.grad(f)(pt)
     self.assertIsInstance(pt, ZeroPoint)
 
+  @parameterized.parameters(1, 2, 3)
+  def test_shape_dtype_struct(self, i):
+    s = api.ShapeDtypeStruct(shape=(i, 2, 3), dtype=np.float32)
+    self.assertEqual(s.shape, (i, 2, 3))
+    self.assertEqual(s.dtype, np.float32)
+    self.assertEqual(s.ndim, 3)
+    self.assertEqual(s.size, i * 2 * 3)
+    self.assertLen(s, i)
+    for f in (str, repr):
+      self.assertEqual(
+          f(s), "ShapeDtypeStruct(shape=({}, 2, 3), dtype=float32)".format(i))
+
+  def test_shape_dtype_struct_scalar(self):
+    s = api.ShapeDtypeStruct(shape=(), dtype=np.float32)
+    self.assertEmpty(s.shape)
+    self.assertEqual(s.size, 1)
+    self.assertEqual(s.ndim, 0)
+    with self.assertRaisesRegex(TypeError, "len[(][)] of unsized object"):
+      _ = len(s)
+
   def test_eval_shape(self):
     def fun(x, y):
       return np.tanh(np.dot(x, y) + 3.)