class_docstr: Helper class for creating sets of parallel function calls.
  
  The purpose of this class is to replace this code:
  
      e = []
      f = []
      for i in range(len(devices)):
        with tf.device(devices[i]):
          e_, f_ = func(a[i], b[i], c)
          e.append(e_)
          f.append(f_)
  
  with this code:
  
      e, f = expert_utils.Parallelism(devices)(func, a, b, c)
class_docstr_tok: ["Helper","class","for","creating","sets","of","parallel","function","calls",".","The","purpose","of","this","class","is","to","replace","this","code",":","e","=","[","]","f","=","[","]","for","i","in","range","(","len","(","devices",")",")",":","with","tf.device","(","devices","[","i","]",")",":","e_",",","f_","=","func","(","a","[","i","]",",","b","[","i","]",",","c",")","e.append","(","e_",")","f.append","(","f_",")","with","this","code",":","e",",","f","=","expert_utils.Parallelism","(","devices",")","(","func",",","a",",","b",",","c",")"]
docstr: A parallel set of function calls (using the specified devices).
  
  Args:
    fn: a function or a list of n functions.
    *args: additional args.  Each arg should either be not a list, or a list
       of length n.
    **kwargs: additional keyword args.  Each arg should either be not a
       list, or a list of length n.
  
  Returns:
    either a single list of length n (if fn does not return a tuple), or a
    tuple of lists of length n (if fn returns a tuple).
docstr_tok: ["A","parallel","set","of","function","calls","(","using","the","specified","devices",")",".","Args",":","fn",":","a","function","or","a","list","of","n","functions",".","*","args",":","additional","args",".","Each","arg","should","either","be","not","a","list",",","or","a","list","of","length","n.","*","*","kwargs",":","additional","keyword","args",".","Each","arg","should","either","be","not","a","list",",","or","a","list","of","length","n.","Returns",":","either","a","single","list","of","length","n","(","if","fn","does","not","return","a","tuple",")",",","or","a","tuple","of","lists","of","length","n","(","if","fn","returns","a","tuple",")","."]
code: def __call__(self, fn, *args, **kwargs):
      """A parallel set of function calls (using the specified devices).
  
      Args:
        fn: a function or a list of n functions.
        *args: additional args.  Each arg should either be not a list, or a list
           of length n.
        **kwargs: additional keyword args.  Each arg should either be not a
           list, or a list of length n.
  
      Returns:
        either a single list of length n (if fn does not return a tuple), or a
        tuple of lists of length n (if fn returns a tuple).
      """
      # Construct lists or args and kwargs for each function.
      if args:
          my_args = transpose_list_of_lists(
              [self._maybe_repeat(arg) for arg in args]
          )
      else:
          my_args = [[] for _ in range(self.n)]
      my_kwargs = [{} for _ in range(self.n)]
      for k, v in six.iteritems(kwargs):
          vals = self._maybe_repeat(v)
          for i in range(self.n):
              my_kwargs[i][k] = vals[i]
  
      # Construct lists of functions.
      fns = self._maybe_repeat(fn)
  
      # Now make the parallel call.
      outputs = []
      cache = {}
      tensor_to_var = {}
      for i in range(self.n):
  
          def daisy_chain_getter(getter, name, *args, **kwargs):
              """Get a variable and cache in a daisy chain."""
              device_var_key = (self._devices[i], name)
              if device_var_key in cache:
                  # if we have the variable on the correct device, return it.
                  return cache[device_var_key]
              if name in cache:
                  # if we have it on a different device, copy it from the last device
                  last_device_v = cache[name]
                  var = tensor_to_var[last_device_v]
                  v = tf.identity(last_device_v)
              else:
                  var = getter(name, *args, **kwargs)
                  v = var.read_value()
  
              # keep track of the original variable
              tensor_to_var[v] = var
              _add_variable_proxy_methods(tensor_to_var[v], v)
              # update the cache
              cache[name] = v
              cache[device_var_key] = v
              return v
  
          # Variable scope will not reset caching_device on reused variables,
          # so we make a custom getter that uses identity to cache the variable.
          # pylint: disable=cell-var-from-loop
          def caching_getter(getter, name, *args, **kwargs):
              """Cache variables on device."""
              key = (self._caching_devices[i], name)
              if key in cache:
                  return cache[key]
  
              v = getter(name, *args, **kwargs)
              with tf.device(self._caching_devices[i]):
                  ret = v.read_value()
              _add_variable_proxy_methods(v, ret)
              cache[key] = ret
              return ret
  
          if self._daisy_chain_variables:
              custom_getter = daisy_chain_getter
          elif self._caching_devices[i]:
              custom_getter = caching_getter
          else:
              custom_getter = None
          # pylint: enable=cell-var-from-loop
          with tf.name_scope("parallel_%d" % i):
              with tf.variable_scope(
                  tf.get_variable_scope()
                  if self._reuse
                  else "parallel_%d" % i,
                  reuse=True if i > 0 and self._reuse else None,
                  caching_device=self._caching_devices[i],
                  custom_getter=custom_getter,
              ):
                  # TODO(noam, epot, avaswani)
                  # Allows for passing no device in case you want to default to the
                  # existing device. This is needed when we put all experts on a single
                  # device, for example in local_moe.
                  if self._devices[i] != DEFAULT_DEV_STRING:
                      with tf.device(self._devices[i]):
                          outputs.append(
                              fns[i](*my_args[i], **my_kwargs[i])
                          )
                  else:
                      outputs.append(
                          fns[i](*my_args[i], **my_kwargs[i])
                      )
      if isinstance(outputs[0], tuple):
          outputs = list(zip(*outputs))
          outputs = tuple([list(o) for o in outputs])
      return outputs
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import functools
  import math
  import six
  from six.moves import range
  from six.moves import zip
  from tensor2tensor.layers import common_layers
  from tensor2tensor.layers.vq_discrete import DiscreteBottleneck
  import tensorflow.compat.v1
label_desc: Use image node_id 2 for calling the Parallelism obj's underlying member method code with example usage: obj.__call__(fn) and returns: outputs, v, ret, cache, cache