code: def device_put(t, device=None):
      return tree.map_structure(
          utils.run_on_device(lambda x: x, device), t
      )
  
file_path: sonnet/sonnet/src/functional/jax.py
imports: import functools
  from sonnet.src.functional import utils
  import tensorflow
  import tree
label_desc: Use image node_id 1 for calling a global function with example usage: device_put(t, device) and returns: tree