code: def _tf_pmap(*args, **kwargs):
      kwargs.pop(
          "donate_argnums", None
      )  # donate_argnums not used in TF
      return tf_np_extensions.pmap(*args, **kwargs)
  
file_path: trax/trax/fastmath/tf.py
imports: import numpy
  import tensorflow.compat.v2
  from trax.shapes import ShapeDtype
  from trax.tf_numpy import extensions
  from trax.tf_numpy import numpy
label_desc: Use image node_id 6 for calling a global function with example usage: _tf_pmap() and returns: tf_np_extensions