Repository: google/jax

Issue:
jit inside pmap inside jit error
Commenting out *either* of the `@jax.jit` decorators works, as does explicitly passing the `count` argument (e.g., `multi_step_pmap(u, 2)`
```python
import jax.numpy as np
import jax
from jax import lax
from functools import partial
from jax.lib import xla_bridge

device_count = xla_bridge.device_count()

@jax.jit
def multi_step(state, count):
  return lax.fori_loop(0, count, lambda i, s: s, state)

@jax.jit
def multi_step_pmap(state, count=2):
  @partial(jax.pmap, axis_name='x')
  def pmapped_multi_step(state):
    return multi_step(state, count)

  return pmapped_multi_step(state)

u = np.ones((1, 100))
u_final = multi_step_pmap(u)
```
Here's the stacktrace:
```
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
<ipython-input-29-4f6d06752943> in <module>()
     20 
     21 u = np.ones((1, 100))
---> 22 u_final = multi_step_pmap(u)

9 frames
/usr/local/lib/python3.6/dist-packages/jax/api.py in f_jitted(*args, **kwargs)
    129     _check_args(args_flat)
    130     flat_fun, out_tree = flatten_fun_leafout(f, in_tree)
--> 131     out = xla.xla_call(flat_fun, *args_flat, device_values=device_values)
    132     return out if out_tree() is leaf else tree_unflatten(out_tree(), out)
    133 

/usr/local/lib/python3.6/dist-packages/jax/core.py in call_bind(primitive, f, *args, **params)
    661   if top_trace is None:
    662     with new_sublevel():
--> 663       ans = primitive.impl(f, *args, **params)
    664   else:
    665     tracers = map(top_trace.full_raise, args)

/usr/local/lib/python3.6/dist-packages/jax/interpreters/xla.py in _xla_call_impl(fun, *args, **params)
    660 def _xla_call_impl(fun, *args, **params):
    661   device_values = FLAGS.jax_device_values and params.pop('device_values')
--> 662   compiled_fun = _xla_callable(fun, device_values, *map(abstractify, args))
    663   try:
    664     return compiled_fun(*args)

/usr/local/lib/python3.6/dist-packages/jax/linear_util.py in memoized_fun(f, *args)
    203 
    204   def memoized_fun(f, *args):
--> 205     ans, f_prev = memoized_fun_body(f, args)
    206     if id(f_prev) != id(f):
    207       f.populate_stores(f_prev)

/usr/local/lib/python3.6/dist-packages/jax/linear_util.py in memoized_fun_body(f, args)
    200   @fastcache.clru_cache(maxsize=max_size)
    201   def memoized_fun_body(f, args):
--> 202     return call(f, *args), f
    203 
    204   def memoized_fun(f, *args):

/usr/local/lib/python3.6/dist-packages/jax/interpreters/xla.py in _xla_callable(fun, device_values, *abstract_args)
    675     assert not env  # no subtraces here (though cond might eventually need them)
    676     axis_env = AxisEnv(jaxpr_replicas(jaxpr), [], [])
--> 677     compiled, result_shape = _compile_jaxpr(jaxpr, axis_env, consts, *abstract_args)
    678     del master, consts, jaxpr, env
    679   if device_values:

/usr/local/lib/python3.6/dist-packages/jax/interpreters/xla.py in _compile_jaxpr(jaxpr, axis_env, const_vals, *abstract_args)
    218     raise ValueErrr(msg.format(axis_env.nreps, xb.device_count()))
    219   arg_shapes = list(map(xla_shape, abstract_args))
--> 220   built_c = _jaxpr_computation(jaxpr, axis_env, const_vals, (), *arg_shapes)
    221   result_shape = xla_shape_to_result_shape(built_c.GetReturnValueShape())
    222   return built_c.Compile(arg_shapes, xb.get_compile_options(axis_env.nreps),

/usr/local/lib/python3.6/dist-packages/jax/interpreters/xla.py in _jaxpr_computation(jaxpr, axis_env, const_vals, freevar_shapes, *arg_shapes)
    265       if isinstance(val, DeviceArray):
    266         val.copy_to_host_async()
--> 267     _map(write, jaxpr.constvars, map(c.Constant, const_vals))
    268     _map(write, jaxpr.freevars, map(c.ParameterWithShape, freevar_shapes))
    269   else:

/usr/local/lib/python3.6/dist-packages/jax/interpreters/xla.py in _map(f, *xs)
    307 
    308 def _map(f, *xs):
--> 309   return tuple(map(f, *xs))
    310 
    311 def xla_destructure(c, ans):

/usr/local/lib/python3.6/dist-packages/jax/lib/xla_bridge.py in Constant(self, py_val, canonicalize_types)
    273       return _constant_handlers[py_type](self, py_val, canonicalize_types)
    274     else:
--> 275       raise TypeError("No constant handler for type: {}".format(py_type))
    276 
    277   # TODO(mattjj): remove when CrossReplicaSum is added to XLA:CPU

TypeError: No constant handler for type: <class 'jax.interpreters.pxla.ShardedDeviceTuple'>
```


Discussion:

Resolving patch:
diff --git a/jax/interpreters/pxla.py b/jax/interpreters/pxla.py
--- a/jax/interpreters/pxla.py
+++ b/jax/interpreters/pxla.py
@@ -399,6 +399,8 @@ def tuple_element_handler(axis_size, aval):
 xla.canonicalize_dtype_handlers[ShardedDeviceTuple] = \
     xla.canonicalize_dtype_handlers[xla.DeviceTuple]
 
+xb.register_constant_handler(ShardedDeviceTuple, xla._device_tuple_constant_handler)
+
 
 class ShardedDeviceArray(ShardedDeviceValue, xla.DeviceArray):
   """A ShardedDeviceArray is an ndarray sharded across devices.


Test patch:
diff --git a/tests/api_test.py b/tests/api_test.py
--- a/tests/api_test.py
+++ b/tests/api_test.py
@@ -17,6 +17,7 @@
 from __future__ import print_function
 
 import collections
+from functools import partial
 
 from absl.testing import absltest
 import numpy as onp
@@ -29,6 +30,7 @@
 from jax.interpreters import ad
 from jax.interpreters.xla import DeviceArray, DeviceTuple
 from jax.abstract_arrays import concretization_err_msg
+from jax.lib import xla_bridge as xb
 from jax import test_util as jtu
 
 from jax.config import config
@@ -899,5 +901,26 @@ def test_jit_of_noncallable(self):
     jtu.check_raises_regexp(lambda: api.jit(3), TypeError,
                             "Expected a callable value.*")
 
+  def test_issue_1062(self):
+    # code from https://github.com/google/jax/issues/1062 @shoyer
+    # this tests, among other things, whether ShardedDeviceTuple constants work
+    device_count = xb.device_count()
+
+    @jit
+    def multi_step(state, count):
+      return lax.fori_loop(0, count, lambda i, s: s, state)
+
+    @jit
+    def multi_step_pmap(state, count=2):
+      @partial(api.pmap, axis_name='x')
+      def pmapped_multi_step(state):
+        return multi_step(state, count)
+
+      return pmapped_multi_step(state)
+
+    u = np.ones((device_count, 100))
+    u_final = multi_step_pmap(u)  # doesn't crash
+
+
 if __name__ == '__main__':
   absltest.main()