Mixed Device Placement in Scan

[ntenenz]

On a GPU-equipped node, I run the following code:

import jax

def fib(c, _):
    next_num = c[0] + c[1]
    return (c[1], next_num), jax.device_put(next_num, jax.devices("cpu")[0])

out = jax.lax.scan(fib, (0, 1), None, 5)

print(out[0][1].device())  # GPU (as expected)
print(out[1].device())     # GPU (not ideal)

It appears scan respects the device of the operation and not the device of the array itself when merging on Y. Is it possible to make it do so or should I resort to a vanilla loop in python?

[YouJiacheng]

scan body will be compiled and only can run on single device IIUC.
FYI: https://jax.readthedocs.io/en/latest/jax.experimental.host_callback.html and https://github.com/mpi4jax/mpi4jax can help you transfer data between devices(and even hosts, using mpi4jax).

[ntenenz]

The use case is a bit similar to the training loop of a neural network (i.e., iterative and computationally expensive -- needs GPU acceleration), however, I'd like to be able to log the history of something akin to gradients (i.e., large amounts of data) throughout the course of training, thus necessitating the movement of data off device.

Are you suggesting an id_tap to pull the data from the device is the way to go?