Repository: google/jax

Issue:
add a shape inference function to api.py
A couple users (including @lukaszkaiser) have requested a way to do shape inference. Maybe we can provide something like `eval_shape` (taking example arguments, as here, or maybe shape tuples):

```python
from __future__ import print_function

# Look at all these internal imports! Not for the faint of heart.
from jax.abstract_arrays import ShapedArray, raise_to_shaped
from jax.api_util import pytree_to_jaxtupletree, pytree_fun_to_jaxtupletree_fun
from jax.core import get_aval
from jax.interpreters import partial_eval as pe
from jax.util import unzip2
import jax.linear_util as lu
from jax.tree_util import build_tree


def abstractify(x):
  return raise_to_shaped(get_aval(x))

def eval_shape(f, *args):
  jaxtuple_args, in_trees = unzip2(map(pytree_to_jaxtupletree, args))
  f, out_tree = pytree_fun_to_jaxtupletree_fun(lu.wrap_init(f), in_trees)
  out = pe.abstract_eval_fun(f.call_wrapped, *map(abstractify, jaxtuple_args))
  return build_tree(out_tree(), out)


###

import jax.numpy as np

def fun(x, y):
  return np.tanh(np.dot(x, y) + 3.)

x = np.ones((2, 3))
y = np.ones((3, 4))
out = eval_shape(fun, x, y)
print(out)


def fun1():
  x = np.ones((2, 3))
  y = np.ones((3, 4))
  return np.tanh(np.dot(x, y) + 3.)
out = eval_shape(fun1)
print(out)


def fun2(x, y):
  a, b = x
  return a + b + y

x = (np.ones(2), np.ones(2))
y = 3.
out = eval_shape(fun2, x, y)
print(out)


def fun3(x, y):
  return x[0] + x[1] + y

x = (np.ones(2), np.ones(2))
y = 3.
out = eval_shape(fun3, x, y)
print(out)


def fun4(x, y):
  return {'hi': x[0] + x[1] + y}

x = (np.ones(2), np.ones(2))
y = 3.
out = eval_shape(fun4, x, y)
print(out)
```


Discussion:
Yes, having a function like this would be amazing! But, more precisely, we'd like to have a function:

eval_shape(fun, *args_shapes)
  return shape of f(ran on rags with shapes as in args_shapes)

2 things that the above code still doesn't do
* using shapes and types instead of values (should be simple)
* removing the need for type of input
  (not super important, but would be very clean; is that technically possible though?)

One more bug to check would be with indexing:

def fun3(x, y):
  a = x[0][0]
  b = x[0][1]
  c = x[1]
  return a + b + c + y

But in general it's a great idea!

Thanks Matt!!
Thanks for catching these bugs, Lukasz. Just updated the code in the OP to handle the `fun3` case. I think it might be getting close now...

Resolving patch:
diff --git a/jax/api.py b/jax/api.py
--- a/jax/api.py
+++ b/jax/api.py
@@ -1062,3 +1062,59 @@ def graphviz_maker(*args, **kwargs):
 
   graphviz_maker.__name__ = "make_graphviz({})".format(graphviz_maker.__name__)
   return graphviz_maker
+
+
+def eval_shape(fun, *args, **kwargs):
+  """Compute the shape of ``fun(*args, **kwargs)`` without incurring 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(np.shape, out)
+
+  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 ``eval_shape`` can also catch shape errors, and will raise same shape
+  errors as evaluating ``fun(*args, **kwargs)``.
+
+  Args:
+    *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:
+
+  >>> f = lambda A, x: np.tanh(np.dot(A, x))
+  >>> class MyArgArray(object):
+  ...   def __init__(self, shape, dtype):
+  ...     self.shape = shape
+  ...     self.dtype = dtype
+  ...
+  >>> A = MyArgArray((2000, 3000), np.float32)
+  >>> x = MyArgArray((3000, 1000), np.float32)
+  >>> out_shape = jax.eval_shape(f, A, x)  # no FLOPs performed
+  >>> print(out_shape)
+  (2000, 1000)
+  """
+  def abstractify(x):
+    if type(x) is core.JaxTuple:
+      return core.AbstractTuple(map(abstractify, x))
+    else:
+      return ShapedArray(onp.shape(x), onp.result_type(x))
+
+  jax_args, in_trees = unzip2(map(pytree_to_jaxtupletree, args))
+  jax_kwargs, kwargs_tree = pytree_to_jaxtupletree(kwargs)
+  f, out_tree = pytree_fun_to_jaxtupletree_fun2(lu.wrap_init(fun), kwargs_tree, in_trees)
+  abstract_args = map(abstractify, (jax_kwargs,) + tuple(jax_args))
+  out = pe.abstract_eval_fun(f.call_wrapped, *abstract_args)
+  return tree_map(onp.shape, build_tree(out_tree(), out))


Test patch:
diff --git a/tests/api_test.py b/tests/api_test.py
--- a/tests/api_test.py
+++ b/tests/api_test.py
@@ -593,5 +593,82 @@ def f(pt):
     f_jit = api.jit(f)
     self.assertAllClose(f(pt), f_jit(pt), check_dtypes=False)
 
+  def test_eval_shape(self):
+    def fun(x, y):
+      return np.tanh(np.dot(x, y) + 3.)
+
+    x = np.ones((2, 3))
+    y = np.ones((3, 4))
+    out_shape = api.eval_shape(fun, x, y)
+
+    self.assertEqual(out_shape, (2, 4))
+
+  def test_eval_shape_constants(self):
+    def fun():
+      x = np.ones((2, 3))
+      y = np.ones((3, 4))
+      return np.tanh(np.dot(x, y) + 3.)
+
+    out_shape = api.eval_shape(fun)
+
+    self.assertEqual(out_shape, (2, 4))
+
+  def test_eval_shape_tuple_unpacking(self):
+    def fun(x, y):
+      a, b = x
+      return a + b + y
+
+    x = (np.ones(2), np.ones(2))
+    y = 3.
+    out_shape = api.eval_shape(fun, x, y)
+
+    self.assertEqual(out_shape, (2,))
+
+  def test_eval_shape_tuple_itemgetting(self):
+    def fun(x, y):
+      return x[0] + x[1] + y
+
+    x = (np.ones(2), np.ones(2))
+    y = 3.
+    out_shape = api.eval_shape(fun, x, y)
+
+    self.assertEqual(out_shape, (2,))
+
+  def test_eval_shape_output_dict(self):
+    def fun(x, y):
+      return {'hi': x[0] + x[1] + y}
+
+    x = (np.ones(2), np.ones(2))
+    y = 3.
+    out_shape = api.eval_shape(fun, x, y)
+
+    self.assertEqual(out_shape, {'hi': (2,)})
+
+  def test_eval_shape_shape_error(self):
+    def fun(x, y):
+      return np.tanh(np.dot(x, y) + 3.)
+
+    x = np.ones((3, 3))
+    y = np.ones((4, 4))
+
+    self.assertRaises(TypeError, lambda: api.eval_shape(fun, x, y))
+
+  def test_eval_shape_duck_typing(self):
+    def fun(A, b, x):
+      return np.dot(A, x) + b
+
+    class MyArgArray(object):
+      def __init__(self, shape, dtype):
+        self.shape = shape
+        self.dtype = dtype
+
+    A = MyArgArray((3, 4), np.float32)
+    b = MyArgArray((5,), np.float32)
+    x = MyArgArray((4, 5), np.float32)
+    out_shape = api.eval_shape(fun, A, b, x)
+
+    self.assertEqual(out_shape, (3, 5))
+
+
 if __name__ == '__main__':
   absltest.main()