Using jax.device_get

[8bitmp3]

Hello w👋 I'm working on a notebook for Flax (@avital) and I noticed jax.device_get doesn't appear to be mentioned in the docs (https://jax.readthedocs.io/en/latest/jax.html) — or maybe I'm not looking in the right place. I've found jax.device_put. Do you think it'd be a good idea to document it?

Use case example:

def evaluate_model(model, test_dataset):
  metrics = eval_step(model, test_dataset)
  # Retrieve stuff from device/pytree registry (?)
  metrics = jax.device_get(fancy_metrics)
  eval_summary = jax.tree_map(lambda x: x.item(), fancy_metrics)
  return eval_summary['test_loss'], eval_summary['test_accuracy']

jax.device_get:

def device_get(x):
  for y in tree_leaves(x):
    try:
      y.copy_to_host_async()
    except AttributeError:
      pass
  return tree_map(_device_get, x)