Cannot grad when parameters are on multiple devices

[ayaka14732]

I am trying to implement pipeline parallelism.

In my model there are four parameters. W1, b1 are on TPU_0, while W2, b2 are on TPU_1.

However, when I use value_and_grad it raises an error.

import jax
import jax.nn as nn
import jax.numpy as np
import jax.random as rand
import numpy as onp

devices = jax.devices()

input_dim: int = 8
output_dim: int = 6
hidden1_dim: int = 4
batch_size: int = 3

def init_params():
    xavier_normal_init = nn.initializers.xavier_normal()
    normal_init = nn.initializers.normal()
    subkey = rand.PRNGKey(hash("let's do it randomly"))  # do not need randomness
    return {
        'W1': jax.device_put(xavier_normal_init(subkey, (input_dim, hidden1_dim)), devices[0]),
        'b1': jax.device_put(normal_init(subkey, (hidden1_dim,)), devices[0]),
        'W2': jax.device_put(xavier_normal_init(subkey, (hidden1_dim, output_dim)), devices[1]),
        'b2': jax.device_put(normal_init(subkey, (output_dim,)), devices[1]),
    }

@jax.value_and_grad
@jax.jit
def forward(params, x):
    x = jax.device_put(nn.relu(np.dot(x, params['W1']) + params['b1']), devices[1])
    x = nn.relu(np.dot(x, params['W2']) + params['b2'])
    x = np.sum(x)
    return x

def main():
    params = init_params()
    x = jax.device_put(onp.ones((batch_size, input_dim)), devices[0])
    value, grad = forward(params, x)
    print(value, grad)

if __name__ == '__main__':
    main()

Traceback (most recent call last):
  File "/home/ayaka/venv38/lib/python3.8/site-packages/jax/_src/dispatch.py", line 116, in _device_from_arg_devices
    device, = {d for d in devices if d is not None} or (None,)
ValueError: too many values to unpack (expected 1)

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "test_grad.py", line 40, in <module>
    main()
  File "test_grad.py", line 36, in main
    value, grad = forward(params, x)
  File "/home/ayaka/venv38/lib/python3.8/site-packages/jax/_src/traceback_util.py", line 162, in reraise_with_filtered_traceback
    return fun(*args, **kwargs)
  File "/home/ayaka/venv38/lib/python3.8/site-packages/jax/_src/api.py", line 1058, in value_and_grad_f
    ans, vjp_py = _vjp(f_partial, *dyn_args, reduce_axes=reduce_axes)
  File "/home/ayaka/venv38/lib/python3.8/site-packages/jax/_src/api.py", line 2466, in _vjp
    out_primal, out_vjp = ad.vjp(
  File "/home/ayaka/venv38/lib/python3.8/site-packages/jax/interpreters/ad.py", line 116, in vjp
    out_primals, pvals, jaxpr, consts = linearize(traceable, *primals)
  File "/home/ayaka/venv38/lib/python3.8/site-packages/jax/interpreters/ad.py", line 103, in linearize
    jaxpr, out_pvals, consts = pe.trace_to_jaxpr(jvpfun_flat, in_pvals)
  File "/home/ayaka/venv38/lib/python3.8/site-packages/jax/_src/profiler.py", line 206, in wrapper
    return func(*args, **kwargs)
  File "/home/ayaka/venv38/lib/python3.8/site-packages/jax/interpreters/partial_eval.py", line 515, in trace_to_jaxpr
    jaxpr, (out_pvals, consts, env) = fun.call_wrapped(pvals)
  File "/home/ayaka/venv38/lib/python3.8/site-packages/jax/linear_util.py", line 166, in call_wrapped
    ans = self.f(*args, **dict(self.params, **kwargs))
  File "/home/ayaka/venv38/lib/python3.8/site-packages/jax/_src/traceback_util.py", line 162, in reraise_with_filtered_traceback
    return fun(*args, **kwargs)
  File "/home/ayaka/venv38/lib/python3.8/site-packages/jax/_src/api.py", line 424, in cache_miss
    out_flat = xla.xla_call(
  File "/home/ayaka/venv38/lib/python3.8/site-packages/jax/core.py", line 1661, in bind
    return call_bind(self, fun, *args, **params)
  File "/home/ayaka/venv38/lib/python3.8/site-packages/jax/core.py", line 1652, in call_bind
    outs = primitive.process(top_trace, fun, tracers, params)
  File "/home/ayaka/venv38/lib/python3.8/site-packages/jax/core.py", line 1664, in process
    return trace.process_call(self, fun, tracers, params)
  File "/home/ayaka/venv38/lib/python3.8/site-packages/jax/interpreters/ad.py", line 323, in process_call
    result = call_primitive.bind(f_jvp, *primals, *nonzero_tangents, **new_params)
  File "/home/ayaka/venv38/lib/python3.8/site-packages/jax/core.py", line 1661, in bind
    return call_bind(self, fun, *args, **params)
  File "/home/ayaka/venv38/lib/python3.8/site-packages/jax/core.py", line 1652, in call_bind
    outs = primitive.process(top_trace, fun, tracers, params)
  File "/home/ayaka/venv38/lib/python3.8/site-packages/jax/core.py", line 1664, in process
    return trace.process_call(self, fun, tracers, params)
  File "/home/ayaka/venv38/lib/python3.8/site-packages/jax/interpreters/partial_eval.py", line 203, in process_call
    jaxpr, out_pvals, consts, env_tracers = self.partial_eval(
  File "/home/ayaka/venv38/lib/python3.8/site-packages/jax/interpreters/partial_eval.py", line 312, in partial_eval
    out_flat, (out_avals, jaxpr, env) = app(f, *in_consts), aux()
  File "/home/ayaka/venv38/lib/python3.8/site-packages/jax/core.py", line 1661, in bind
    return call_bind(self, fun, *args, **params)
  File "/home/ayaka/venv38/lib/python3.8/site-packages/jax/core.py", line 1652, in call_bind
    outs = primitive.process(top_trace, fun, tracers, params)
  File "/home/ayaka/venv38/lib/python3.8/site-packages/jax/core.py", line 1664, in process
    return trace.process_call(self, fun, tracers, params)
  File "/home/ayaka/venv38/lib/python3.8/site-packages/jax/core.py", line 633, in process_call
    return primitive.impl(f, *tracers, **params)
  File "/home/ayaka/venv38/lib/python3.8/site-packages/jax/_src/dispatch.py", line 128, in _xla_call_impl
    compiled_fun = _xla_callable(fun, device, backend, name, donated_invars,
  File "/home/ayaka/venv38/lib/python3.8/site-packages/jax/linear_util.py", line 263, in memoized_fun
    ans = call(fun, *args)
  File "/home/ayaka/venv38/lib/python3.8/site-packages/jax/_src/dispatch.py", line 155, in _xla_callable_uncached
    return lower_xla_callable(fun, device, backend, name, donated_invars,
  File "/home/ayaka/venv38/lib/python3.8/site-packages/jax/_src/profiler.py", line 206, in wrapper
    return func(*args, **kwargs)
  File "/home/ayaka/venv38/lib/python3.8/site-packages/jax/_src/dispatch.py", line 184, in lower_xla_callable
    device = _xla_callable_device(nreps, backend, device, arg_devices)
  File "/home/ayaka/venv38/lib/python3.8/site-packages/jax/_src/dispatch.py", line 350, in _xla_callable_device
    return _device_from_arg_devices(arg_devices)
  File "/home/ayaka/venv38/lib/python3.8/site-packages/jax/_src/dispatch.py", line 120, in _device_from_arg_devices
    raise ValueError(msg.format(", ".join(map(str, devices)))) from err
jax._src.traceback_util.UnfilteredStackTrace: ValueError: primitive arguments must be colocated on the same device, got TPU_0(process=0,(0,0,0,0)), TPU_1(process=0,(0,0,0,1)), TPU_0(process=0,(0,0,0,0)), TPU_1(process=0,(0,0,0,1)), TPU_0(process=0,(0,0,0,0))

The stack trace below excludes JAX-internal frames.
The preceding is the original exception that occurred, unmodified.

--------------------

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "test_grad.py", line 40, in <module>
    main()
  File "test_grad.py", line 36, in main
    value, grad = forward(params, x)
  File "/home/ayaka/venv38/lib/python3.8/site-packages/jax/_src/dispatch.py", line 120, in _device_from_arg_devices
    raise ValueError(msg.format(", ".join(map(str, devices)))) from err
ValueError: primitive arguments must be colocated on the same device, got TPU_0(process=0,(0,0,0,0)), TPU_1(process=0,(0,0,0,1)), TPU_0(process=0,(0,0,0,0)), TPU_1(process=0,(0,0,0,1)), TPU_0(process=0,(0,0,0,0))