code: def device_get(t):
      return tree.map_structure(lambda x: x.numpy(), 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 2 for calling a global function with example usage: device_get(t) and returns: tree