Implement size/ndim/__len__/repr/str/eq/hash for ShapeDtypeStruct. (jax-ml#2206)

Author: tomhennigan

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.
 

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.)