Repository: google/jax

Issue:
jit(..., backend='cpu') and the new device placement don't go together
After #1884 this code snippet fails:
```
    def get_arr(scale):
      return scale + np.ones((2, 2))

    get_arr = jit(get_arr, backend='cpu')
    x = get_arr(0.1)
    x / x.shape[0]
    x + np.ones_like(x)
    x + np.eye(2)
```
with errors like 
```
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
<ipython-input-20-ca9775de72ea> in <module>()
----> 1 x / x.shape[0]

4 frames
google3/third_party/py/jax/interpreters/xla.py in xla_primitive_callable(prim, *arg_specs, **params)
    169     msg = "primitive arguments must be colocated on the same device, got {}"
    170     names = ("{}({})".format(d[0].__name__, d[1]) for d in devices if d is not None)
--> 171     raise ValueError(msg.format(", ".join(names)))
    172   else:
    173     all_devices = it.chain(xb.devices(), xb.devices('cpu'))

ValueError: primitive arguments must be colocated on the same device, got CpuDevice(0), GpuDevice(0)
```
Note that it works file if instead of `jit(..., backend='cpu')` I use `jax.device_get`.

Below is the modified `tests/multi_device_test.py` to test for this case.
```
# Copyright 2019 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import os
from unittest import SkipTest

from absl.testing import absltest
import numpy as onp
from jax.api import jit
import jax
import jax.numpy as np
from jax import lax
from jax import test_util as jtu
from jax.lib import xla_bridge
from jax.interpreters import xla

from jax.config import config
config.parse_flags_with_absl()

prev_xla_flags = None


# Run all tests with 8 CPU devices.
def setUpModule():
  global prev_xla_flags
  prev_xla_flags = os.getenv("XLA_FLAGS")
  flags_str = prev_xla_flags or ""
  # Don't override user-specified device count, or other XLA flags.
  if "xla_force_host_platform_device_count" not in flags_str:
    os.environ["XLA_FLAGS"] = (flags_str +
                               " --xla_force_host_platform_device_count=8")
  # Clear any cached backends so new CPU backend will pick up the env var.
  xla_bridge.get_backend.cache_clear()

# Reset to previous configuration in case other test modules will be run.
def tearDownModule():
  if prev_xla_flags is None:
    del os.environ["XLA_FLAGS"]
  else:
    os.environ["XLA_FLAGS"] = prev_xla_flags
  xla_bridge.get_backend.cache_clear()


class MultiDeviceTest(jtu.JaxTestCase):

  def test_computation_follows_data(self):
    if len(jax.devices()) < 2:
      raise SkipTest("test requires multiple devices")

    x = jax.device_put(1, jax.devices()[0])
    y = jax.device_put(2, jax.devices()[0])
    z = x + y
    self.assertEqual(z, 3)
    self.assertEqual(z.device_buffer.device(), jax.devices()[0])

    x = jax.device_put(1, jax.devices()[1])
    y = jax.device_put(2, jax.devices()[1])
    z = x + y
    self.assertEqual(z, 3)
    self.assertEqual(z.device_buffer.device(), jax.devices()[1])

    x = jax.device_put(1, jax.devices()[1])
    y = 4
    z = x + y
    self.assertEqual(z, 5)
    self.assertEqual(z.device_buffer.device(), jax.devices()[1])

    x = jax.device_put(1, jax.devices()[1])
    y = np.ones(3)
    z = x + y
    self.assertAllClose(z, 1 + onp.ones(3), check_dtypes=False)
    self.assertEqual(z.device_buffer.device(), jax.devices()[1])

    x = jax.device_put(1, jax.devices()[0])
    y = jax.device_put(2, jax.devices()[1])
    self.assertRaisesRegex(
        ValueError,
        "primitive arguments must be colocated on the same device",
        lambda: x + y)

    x = jax.device_put(1, jax.devices()[1])
    y = x.reshape((1, 1))
    self.assertEqual(y.device_buffer.device(), jax.devices()[1])

  def test_primitive_compilation_cache(self):
    if len(jax.devices()) < 2:
      raise SkipTest("test requires multiple devices")

    primitive_computation = xla.primitive_computation
    xla.xla_primitive_callable.cache_clear()  # clear op-by-op cache

    count = [0]
    def primitive_computation_and_count(*args, **kwargs):
      count[0] += 1
      return primitive_computation(*args, **kwargs)

    x = jax.device_put(1, jax.devices()[1])

    try:
      xla.primitive_computation = primitive_computation_and_count
      y = lax.add(x, x)
      z = lax.add(y, y)
    finally:
      xla.primitive_computation = primitive_computation

    self.assertEqual(count[0], 1)
    self.assertEqual(y.device_buffer.device(), jax.devices()[1])
    self.assertEqual(z.device_buffer.device(), jax.devices()[1])

  def test_jit_cpu(self):
    def get_arr(scale):
      return scale + np.ones((2, 2))

    get_arr = jit(get_arr, backend='cpu')
    x = get_arr(0.1)
    x / x.shape[0]
    x + np.ones_like(x)
    x + np.eye(2)


if __name__ == '__main__':
  absltest.main()
```
jit(..., backend='cpu') and the new device placement don't go together
After #1884 this code snippet fails:
```
    def get_arr(scale):
      return scale + np.ones((2, 2))

    get_arr = jit(get_arr, backend='cpu')
    x = get_arr(0.1)
    x / x.shape[0]
    x + np.ones_like(x)
    x + np.eye(2)
```
with errors like 
```
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
<ipython-input-20-ca9775de72ea> in <module>()
----> 1 x / x.shape[0]

4 frames
google3/third_party/py/jax/interpreters/xla.py in xla_primitive_callable(prim, *arg_specs, **params)
    169     msg = "primitive arguments must be colocated on the same device, got {}"
    170     names = ("{}({})".format(d[0].__name__, d[1]) for d in devices if d is not None)
--> 171     raise ValueError(msg.format(", ".join(names)))
    172   else:
    173     all_devices = it.chain(xb.devices(), xb.devices('cpu'))

ValueError: primitive arguments must be colocated on the same device, got CpuDevice(0), GpuDevice(0)
```
Note that it works file if instead of `jit(..., backend='cpu')` I use `jax.device_get`.

Below is the modified `tests/multi_device_test.py` to test for this case.
```
# Copyright 2019 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import os
from unittest import SkipTest

from absl.testing import absltest
import numpy as onp
from jax.api import jit
import jax
import jax.numpy as np
from jax import lax
from jax import test_util as jtu
from jax.lib import xla_bridge
from jax.interpreters import xla

from jax.config import config
config.parse_flags_with_absl()

prev_xla_flags = None


# Run all tests with 8 CPU devices.
def setUpModule():
  global prev_xla_flags
  prev_xla_flags = os.getenv("XLA_FLAGS")
  flags_str = prev_xla_flags or ""
  # Don't override user-specified device count, or other XLA flags.
  if "xla_force_host_platform_device_count" not in flags_str:
    os.environ["XLA_FLAGS"] = (flags_str +
                               " --xla_force_host_platform_device_count=8")
  # Clear any cached backends so new CPU backend will pick up the env var.
  xla_bridge.get_backend.cache_clear()

# Reset to previous configuration in case other test modules will be run.
def tearDownModule():
  if prev_xla_flags is None:
    del os.environ["XLA_FLAGS"]
  else:
    os.environ["XLA_FLAGS"] = prev_xla_flags
  xla_bridge.get_backend.cache_clear()


class MultiDeviceTest(jtu.JaxTestCase):

  def test_computation_follows_data(self):
    if len(jax.devices()) < 2:
      raise SkipTest("test requires multiple devices")

    x = jax.device_put(1, jax.devices()[0])
    y = jax.device_put(2, jax.devices()[0])
    z = x + y
    self.assertEqual(z, 3)
    self.assertEqual(z.device_buffer.device(), jax.devices()[0])

    x = jax.device_put(1, jax.devices()[1])
    y = jax.device_put(2, jax.devices()[1])
    z = x + y
    self.assertEqual(z, 3)
    self.assertEqual(z.device_buffer.device(), jax.devices()[1])

    x = jax.device_put(1, jax.devices()[1])
    y = 4
    z = x + y
    self.assertEqual(z, 5)
    self.assertEqual(z.device_buffer.device(), jax.devices()[1])

    x = jax.device_put(1, jax.devices()[1])
    y = np.ones(3)
    z = x + y
    self.assertAllClose(z, 1 + onp.ones(3), check_dtypes=False)
    self.assertEqual(z.device_buffer.device(), jax.devices()[1])

    x = jax.device_put(1, jax.devices()[0])
    y = jax.device_put(2, jax.devices()[1])
    self.assertRaisesRegex(
        ValueError,
        "primitive arguments must be colocated on the same device",
        lambda: x + y)

    x = jax.device_put(1, jax.devices()[1])
    y = x.reshape((1, 1))
    self.assertEqual(y.device_buffer.device(), jax.devices()[1])

  def test_primitive_compilation_cache(self):
    if len(jax.devices()) < 2:
      raise SkipTest("test requires multiple devices")

    primitive_computation = xla.primitive_computation
    xla.xla_primitive_callable.cache_clear()  # clear op-by-op cache

    count = [0]
    def primitive_computation_and_count(*args, **kwargs):
      count[0] += 1
      return primitive_computation(*args, **kwargs)

    x = jax.device_put(1, jax.devices()[1])

    try:
      xla.primitive_computation = primitive_computation_and_count
      y = lax.add(x, x)
      z = lax.add(y, y)
    finally:
      xla.primitive_computation = primitive_computation

    self.assertEqual(count[0], 1)
    self.assertEqual(y.device_buffer.device(), jax.devices()[1])
    self.assertEqual(z.device_buffer.device(), jax.devices()[1])

  def test_jit_cpu(self):
    def get_arr(scale):
      return scale + np.ones((2, 2))

    get_arr = jit(get_arr, backend='cpu')
    x = get_arr(0.1)
    x / x.shape[0]
    x + np.ones_like(x)
    x + np.eye(2)


if __name__ == '__main__':
  absltest.main()
```


Discussion:
IIUC the test only fails on a GPU machine, or at least one where there's a backend other than `'cpu'`. Does that sound right?
IIUC the test only fails on a GPU machine, or at least one where there's a backend other than `'cpu'`. Does that sound right?

Resolving patch:
diff --git a/jax/interpreters/pxla.py b/jax/interpreters/pxla.py
--- a/jax/interpreters/pxla.py
+++ b/jax/interpreters/pxla.py
@@ -360,7 +360,7 @@ def __getitem__(self, idx):
       ids = self._ids()
       device_buffer = self.device_buffers[ids[idx]]
       aval = ShapedArray(self.aval.shape[1:], self.aval.dtype)
-      handler = xla.aval_to_result_handler(aval)
+      handler = xla.aval_to_result_handler(None, aval)
       return handler(device_buffer)
     else:
       return super(ShardedDeviceArray, self).__getitem__(idx)
diff --git a/jax/interpreters/xla.py b/jax/interpreters/xla.py
--- a/jax/interpreters/xla.py
+++ b/jax/interpreters/xla.py
@@ -70,14 +70,15 @@ def aval_to_xla_shape(aval):
 xla_shape_handlers[ShapedArray] = lambda a: xc.Shape.array_shape(a.dtype, a.shape)
 xla_shape_handlers[ConcreteArray] = lambda a: xc.Shape.array_shape(a.dtype, a.shape)
 
-def aval_to_result_handler(aval):
+def aval_to_result_handler(device, aval):
   try:
-    return xla_result_handlers[type(aval)](aval)
+    return xla_result_handlers[type(aval)](device, aval)
   except KeyError:
     raise TypeError("No xla_result_handler for type: {}".format(type(aval)))
 xla_result_handlers = {}
-xla_result_handlers[core.AbstractUnit] = lambda _: lambda _: core.unit
-def array_result_handler(aval): return partial(DeviceArray, raise_to_shaped(aval))
+xla_result_handlers[core.AbstractUnit] = lambda _, __: lambda _: core.unit
+def array_result_handler(device, aval):
+  return partial(DeviceArray, raise_to_shaped(aval), device)
 xla_result_handlers[ShapedArray] = array_result_handler
 xla_result_handlers[ConcreteArray] = array_result_handler
 
@@ -147,11 +148,6 @@ def _make_abstract_python_scalar(typ, _):
 
 ### op-by-op execution
 
-def apply_primitive(prim, *args, **params):
-  """Impl rule that compiles and runs a single primitive 'prim' using XLA."""
-  compiled_fun = xla_primitive_callable(prim, *map(arg_spec, args), **params)
-  return compiled_fun(*args)
-
 def arg_spec(x):
   aval = abstractify(x)
   try:
@@ -159,33 +155,51 @@ def arg_spec(x):
   except:
     return aval, None
 
+def apply_primitive(prim, *args, **params):
+  """Impl rule that compiles and runs a single primitive 'prim' using XLA."""
+  compiled_fun = xla_primitive_callable(prim, *map(arg_spec, args), **params)
+  return compiled_fun(*args)
+
 @cache()
 def xla_primitive_callable(prim, *arg_specs, **params):
-  avals, devices = unzip2(arg_specs)
-  # TODO(mattjj): make Device hashable instead of handling pairs here
-  try:
-    device, = set(d for d in devices if d is not None) or (None,)
-  except ValueError:
-    msg = "primitive arguments must be colocated on the same device, got {}"
-    names = ("{}({})".format(d[0].__name__, d[1]) for d in devices if d is not None)
-    raise ValueError(msg.format(", ".join(names)))
-  else:
-    all_devices = it.chain(xb.devices(), xb.devices('cpu'))
-    device = device and next(d for d in all_devices if (type(d), d.id) == device)
+  avals, arg_devices = unzip2(arg_specs)
+  device = _device_from_arg_devices(arg_devices)
   backend = xb.get_device_backend(device)
   aval_out = prim.abstract_eval(*avals, **params)
-  if prim.multiple_results:
-    handlers = tuple(map(aval_to_result_handler, aval_out))
-    handle_result = lambda xs: tuple(h(x) for h, x in zip(handlers, xs.destructure()))
+  if not prim.multiple_results:
+    handle_result = aval_to_result_handler(device, aval_out)
   else:
-    handle_result = aval_to_result_handler(aval_out)
+    handlers = tuple(map(partial(aval_to_result_handler, device), aval_out))
+    handle_result = lambda xs: tuple(h(x) for h, x in zip(handlers, xs.destructure()))
   tuple_args = len(avals) > 100
   built_c = primitive_computation(prim, backend, tuple_args, *avals, **params)
-  options = xb.get_compile_options(device_assignment=(device.id,) if device else None)
+  options = xb.get_compile_options(device_assignment=device and (device.id,))
   compiled = built_c.Compile(compile_options=options, backend=backend)
   return partial(_execute_compiled_primitive, prim, compiled, backend,
                  tuple_args, handle_result)
 
+# TODO(mattjj): make Device instances hashable instead of handling pairs here
+def _device_from_arg_devices(devices):
+  """Given devices of inputs, determine where to perform a computation.
+
+  Args:
+    devices: list where each element is a either a pair consisting of a device
+      class and an int id (representing a Device instance) or a None.
+  Returns:
+    A Device instance or None.
+  Raises:
+    ValueError if input devices are inconsistent.
+  """
+  try:
+    device, = set(d for d in devices if d is not None) or (None,)
+  except ValueError:
+    msg = "primitive arguments must be colocated on the same device, got {}"
+    names = ("{}({})".format(d[0].__name__, d[1]) for d in devices if d is not None)
+    raise ValueError(msg.format(", ".join(names)))
+  else:
+    all_devices = it.chain(xb.devices(), xb.devices('cpu'))
+    return device and next(d for d in all_devices if (type(d), d.id) == device)
+
 @cache()
 def primitive_computation(prim, backend, tuple_args, *avals, **params):
   c = xb.make_computation_builder("primitive_computation_{}".format(prim.name))
@@ -425,7 +439,7 @@ def eqn_collectives(eqn):
 def _xla_call_impl(fun, *args, **params):
   device = params['device']
   backend = params['backend']
-  compiled_fun = _xla_callable(fun, device, backend, *map(abstractify, args))
+  compiled_fun = _xla_callable(fun, device, backend, *map(arg_spec, args))
   try:
     return compiled_fun(*args)
   except FloatingPointError:
@@ -434,33 +448,38 @@ def _xla_call_impl(fun, *args, **params):
     return fun.call_wrapped(*args)  # probably won't return
 
 @lu.cache
-def _xla_callable(fun, device, backend, *abstract_args):
+def _xla_callable(fun, device, backend, *arg_specs):
+  if device is not None and backend is not None:
+    raise ValueError("can't specify both a device and a backend for jit, "
+                     "got device={} and backend={}".format(device, backend))
+
+  abstract_args, arg_devices = unzip2(arg_specs)
   pvals = [pe.PartialVal((aval, core.unit)) for aval in abstract_args]
   with core.new_master(pe.StagingJaxprTrace, True) as master:
     jaxpr, (pvals, consts, env) = pe.trace_to_subjaxpr(fun, master, False).call_wrapped(pvals)
     assert not env  # no subtraces here
     del master, env
   _map(prefetch, it.chain(consts, jaxpr_literals(jaxpr)))
-  result_handlers = tuple(map(_pval_to_result_handler, pvals))
+
+  nreps = jaxpr_replicas(jaxpr)
+  device = _xla_callable_device(nreps, backend, device, arg_devices)
+  result_handlers = tuple(map(partial(_pval_to_result_handler, device), pvals))
 
   # Computations that only produce constants and/or only rearrange their inputs,
   # which are often produced from partial evaluation, don't need compilation,
   # and don't need to force their (potentially lazy) arguments.
   if not jaxpr.eqns:
-    device = _get_device(device, backend)
+    device = device or xb.get_backend(None).get_default_device_assignment(1)[0]
     return partial(_execute_trivial, jaxpr, device, consts, result_handlers)
 
   log_priority = logging.WARNING if FLAGS.jax_log_compiles else logging.DEBUG
   logging.log(log_priority,
               "Compiling {} for args {}.".format(fun.__name__, abstract_args))
 
-  nreps = jaxpr_replicas(jaxpr)
   if nreps > xb.device_count(backend):
     msg = ("compiling computation that requires {} replicas, but only {} XLA "
             "devices are available")
     raise ValueError(msg.format(nreps, xb.device_count(backend)))
-  axis_env = AxisEnv(nreps, [], [])
-
   if xb.host_count() > 1 and (nreps > 1 or jaxpr_has_pmap(jaxpr)):
     raise NotImplementedError(
         "jit of multi-host pmap not implemented (and jit-of-pmap can cause "
@@ -471,11 +490,10 @@ def _xla_callable(fun, device, backend, *abstract_args):
   c = xb.make_computation_builder("jit_{}".format(fun.__name__))
   xla_consts = _map(c.Constant, consts)
   xla_args = _xla_callable_args(c, abstract_args, tuple_args)
-  out_nodes = jaxpr_subcomp(c, jaxpr, backend, axis_env, xla_consts, (), *xla_args)
+  out_nodes = jaxpr_subcomp(c, jaxpr, backend, AxisEnv(nreps, [], []),
+                            xla_consts, (), *xla_args)
   built = c.Build(c.Tuple(*out_nodes))
 
-  if device is not None and nreps > 1:
-    raise ValueError("can't specify device assignment for jit-of-pmap")
   options = xb.get_compile_options(
       num_replicas=nreps, device_assignment=(device.id,) if device else None)
   compiled = built.Compile(compile_options=options, backend=xb.get_backend(backend))
@@ -485,6 +503,22 @@ def _xla_callable(fun, device, backend, *abstract_args):
   else:
     return partial(_execute_replicated, compiled, backend, result_handlers, tuple_args)
 
+def _xla_callable_device(nreps, backend, device, arg_devices):
+  if nreps > 1:
+    if device is not None or backend is not None:
+      raise ValueError("can't specify device or backend for jit-of-pmap, "
+                       "got device={} and backend={}".format(device, backend))
+    return None
+  else:
+    if device is None and backend is None:
+      return _device_from_arg_devices(arg_devices)
+    elif device is not None and backend is None:
+      return device
+    elif device is None and backend is not None:
+      return xb.get_backend(backend).get_default_device_assignment(1)[0]
+    else:
+      assert False  # Unreachable given the error check in _xla_callable
+
 def _xla_callable_args(c, avals, tuple_args):
   if not tuple_args:
     xla_args = [c.ParameterWithShape(aval_to_xla_shape(a))
@@ -499,12 +533,12 @@ def _xla_callable_args(c, avals, tuple_args):
     assert next(xla_inputs, None) is None
     return xla_args
 
-def _pval_to_result_handler(pval):
+def _pval_to_result_handler(device, pval):
   pv, const = pval
   if pv is None:
     return lambda _: const
   else:
-    return aval_to_result_handler(pv)
+    return aval_to_result_handler(device, pv)
 
 def _execute_compiled(compiled, backend, handlers, tuple_args, *args):
   device, = compiled.local_devices()
@@ -631,7 +665,7 @@ class Token(object): pass
 pytype_aval_mappings[Token] = lambda _: abstract_token
 core.pytype_aval_mappings[Token] = lambda _: abstract_token
 xla_shape_handlers[AbstractToken] = lambda _: xc.Shape.token_shape()
-xla_result_handlers[AbstractToken] = lambda _: lambda _: token
+xla_result_handlers[AbstractToken] = lambda _, __: lambda _: token
 canonicalize_dtype_handlers[Token] = identity
 
 
@@ -671,11 +705,9 @@ class DeviceArray(DeviceValue):
   __slots__ = ["_npy_value", "_device"]
   __array_priority__ = 100
 
-  def __init__(self, aval, device_buffer):
+  def __init__(self, aval, device, device_buffer):
     self.aval = aval
     self.device_buffer = device_buffer
-    # TODO(mattjj): make Device hashable
-    device = device_buffer.device()
     self._device = device and (type(device), device.id)
 
     self._npy_value = None
@@ -835,7 +867,7 @@ def _device_put_impl(x, device=None):
   except TypeError:
     raise TypeError("Argument '{}' of type {} is not a valid JAX type"
                     .format(x, type(x)))
-  handler = aval_to_result_handler(a)
+  handler = aval_to_result_handler(device, a)
   return handler(device_put(x, device))
 
 device_put_p = core.Primitive('device_put')
@@ -887,13 +919,13 @@ def _instantiate_device_constant(const, device=None, backend=None, cutoff=1e6):
   # dispatch an XLA Computation to build the constant on the device if it's
   # large, or alternatively build it on the host and transfer it if it's small
   assert isinstance(const, DeviceConstant)
+  backend = xb.get_backend(device.platform) if device else xb.get_backend(backend)
   if const.size > cutoff:
     c = xb.make_computation_builder("constant_instantiating_computation")
     xla_const = const.constant_handler(c, const)
     device_assignment = (device.id,) if device else None
     opts = xb.get_compile_options(device_assignment=device_assignment)
-    compiled = c.Build(xla_const).Compile((), opts, backend=xb.get_backend(backend))
+    compiled = c.Build(xla_const).Compile((), opts, backend=backend)
     return compiled.Execute(())
   else:
-    return xc.Buffer.from_pyval(onp.asarray(const), device,
-                                backend=xb.get_backend(backend))
+    return xc.Buffer.from_pyval(onp.asarray(const), device, backend=backend)


Test patch:
diff --git a/tests/fft_test.py b/tests/fft_test.py
--- a/tests/fft_test.py
+++ b/tests/fft_test.py
@@ -79,8 +79,7 @@ def testFftn(self, inverse, shape, dtype, axes, rng_factory):
     self._CompileAndCheck(np_fn, args_maker, check_dtypes=True)
     # Test gradient for differentiable types.
     if dtype in inexact_dtypes:
-      # TODO(skye): can we be more precise?
-      tol = 1e-1
+      tol = 0.15  # TODO(skye): can we be more precise?
       jtu.check_grads(np_fn, args_maker(), order=1, atol=tol, rtol=tol)
       jtu.check_grads(np_fn, args_maker(), order=2, atol=tol, rtol=tol)
 
diff --git a/tests/multi_device_test.py b/tests/multi_device_test.py
--- a/tests/multi_device_test.py
+++ b/tests/multi_device_test.py
@@ -62,30 +62,77 @@ def test_computation_follows_data(self):
     if len(jax.devices()) < 2:
       raise SkipTest("test requires multiple devices")
 
+    # computation follows data explicitly placed on device 1
+    x = jax.device_put(1, jax.devices()[1])
+    y = x.reshape((1, 1))
+    self.assertEqual(y.device_buffer.device(), jax.devices()[1])
+    z = y.reshape((1, 1))
+    self.assertEqual(z.device_buffer.device(), jax.devices()[1])
+
+    # multiple arguments explicitly placed on device 0 are compatible
     x = jax.device_put(1, jax.devices()[0])
     y = jax.device_put(2, jax.devices()[0])
     z = x + y
     self.assertEqual(z, 3)
     self.assertEqual(z.device_buffer.device(), jax.devices()[0])
+    w = z + x
+    self.assertEqual(w.device_buffer.device(), jax.devices()[0])
 
+    f = jax.jit(lambda x: x + 1, device=jax.devices()[0])
+    z = f(1) + f(2)
+    self.assertEqual(z, 5)
+    self.assertEqual(z.device_buffer.device(), jax.devices()[0])
+    w = z + z
+    self.assertEqual(z.device_buffer.device(), jax.devices()[0])
+
+    # multiple arguments explicitly placed on device 1 are compatible
     x = jax.device_put(1, jax.devices()[1])
     y = jax.device_put(2, jax.devices()[1])
     z = x + y
     self.assertEqual(z, 3)
     self.assertEqual(z.device_buffer.device(), jax.devices()[1])
+    w = z + x
+    self.assertEqual(z.device_buffer.device(), jax.devices()[1])
 
-    x = jax.device_put(1, jax.devices()[1])
-    y = 4
-    z = x + y
+    f = jax.jit(lambda x: x + 1, device=jax.devices()[1])
+    z = f(1) + f(2)
     self.assertEqual(z, 5)
     self.assertEqual(z.device_buffer.device(), jax.devices()[1])
+    w = z + z
+    self.assertEqual(z.device_buffer.device(), jax.devices()[1])
 
-    x = jax.device_put(1, jax.devices()[1])
-    y = np.ones(3)
-    z = x + y
+    # an argument explicitly placed on one device still works with values that
+    # aren't device-committed (and computaiton follows device-committed values)
+    z = jax.device_put(1., jax.devices()[1]) + 4
+    self.assertEqual(z, 5.)
+    self.assertEqual(z.device_buffer.device(), jax.devices()[1])
+    w = z + 3
+    self.assertEqual(w, 8.)
+    self.assertEqual(w.device_buffer.device(), jax.devices()[1])
+
+    z = jax.device_put(1., jax.devices()[1]) + np.ones(3)
     self.assertAllClose(z, 1 + onp.ones(3), check_dtypes=False)
     self.assertEqual(z.device_buffer.device(), jax.devices()[1])
+    w = z - 3
+    self.assertAllClose(w, 1 + onp.ones(3) - 3, check_dtypes=False)
+    self.assertEqual(w.device_buffer.device(), jax.devices()[1])
+
+    z = jax.device_put(1., jax.devices()[1]) + np.array([1, 2])
+    self.assertAllClose(z, 1 + onp.array([1, 2]), check_dtypes=False)
+    self.assertEqual(z.device_buffer.device(), jax.devices()[1])
+    w = z * 2
+    self.assertAllClose(w, (1 + onp.array([1, 2])) * 2, check_dtypes=False)
+    self.assertEqual(w.device_buffer.device(), jax.devices()[1])
+
+    z = jax.device_put(1., jax.devices()[1]) + jax.device_put(2)
+    self.assertAllClose(z, 3., check_dtypes=False)
+    self.assertEqual(z.device_buffer.device(), jax.devices()[1])
+
+    z = jax.device_put(1., jax.devices()[1]) + jax.jit(lambda x: x + 1)(3)
+    self.assertAllClose(z, 5., check_dtypes=False)
+    self.assertEqual(z.device_buffer.device(), jax.devices()[1])
 
+    # multiple arguments explicitly placed on distinct devices cause errors
     x = jax.device_put(1, jax.devices()[0])
     y = jax.device_put(2, jax.devices()[1])
     self.assertRaisesRegex(
@@ -93,9 +140,12 @@ def test_computation_follows_data(self):
         "primitive arguments must be colocated on the same device",
         lambda: x + y)
 
-    x = jax.device_put(1, jax.devices()[1])
-    y = x.reshape((1, 1))
-    self.assertEqual(y.device_buffer.device(), jax.devices()[1])
+    f = jax.jit(lambda x: x + 1, device=jax.devices()[0])
+    g = jax.jit(lambda x: x + 1, device=jax.devices()[1])
+    self.assertRaisesRegex(
+        ValueError,
+        "primitive arguments must be colocated on the same device",
+        lambda: f(1) + g(2))
 
   def test_primitive_compilation_cache(self):
     if len(jax.devices()) < 2:
diff --git a/tests/multibackend_test.py b/tests/multibackend_test.py
--- a/tests/multibackend_test.py
+++ b/tests/multibackend_test.py
@@ -123,6 +123,22 @@ def fun(x, y):
     self.assertEqual(z.device_buffer.platform(), backend)
     self.assertEqual(w.device_buffer.platform(), backend)
 
+  @jtu.skip_on_devices("cpu")  # test can only fail with non-cpu backends
+  def testJitCpu(self):
+    @partial(api.jit, backend='cpu')
+    def get_arr(scale):
+      return scale + np.ones((2, 2))
+
+    x = get_arr(0.1)
+
+    a = x / x.shape[0]
+    b = x + np.ones_like(x)
+    c = x + np.eye(2)
+
+    self.assertEqual(a.device_buffer.device(), api.devices('cpu')[0])
+    self.assertEqual(b.device_buffer.device(), api.devices('cpu')[0])
+    self.assertEqual(c.device_buffer.device(), api.devices('cpu')[0])
+
 
 if __name__ == "__main__":
   absltest.main()