code: def to_tf_fn(f):
      return lambda *args: f(*to_np(args))
  
imports: import functools
  import itertools
  from absl import flags
  from absl.testing import parameterized
  from jax import lax
  import numpy
  import tensorflow.compat.v2
  from trax.tf_numpy import extensions
  import trax.tf_numpy.numpy
label_desc: Use image node_id 6 for calling a global function with example usage: to_tf_fn(f) and returns: function