code: def _apply_optimizer(data, apply_fn):
      params, per_step_grads = data
      params = [
          tf.Variable(p, name="rank{}".format(len(p.shape)))
          for p in params
      ]
      per_step_grads = tree.map_structure(
          tf.convert_to_tensor, per_step_grads
      )
      param_vals = []
      assert per_step_grads
      for grads in per_step_grads:
          apply_fn(grads, params)
          param_vals.append([p.numpy() for p in params])
      return param_vals
  
imports: import itertools
  from absl.testing import parameterized
  import numpy
  from sonnet.src import base
  from sonnet.src import test_utils
  import tensorflow
  import tree
label_desc: Use image node_id 3 for calling a global function with example usage: _apply_optimizer(data, apply_fn) and returns: param_vals