(f, devices, has_tpu)
This is a helper function to return the pmap impl.

  Args:
    f: a function that takes ndarrays and returns ndarrays.
    devices: a list of strings; the device list.
    has_tpu: boolean; whether `devices` contains TPU devices.

  Returns:
    A function that takes tensors and returns tensors.
This is a helper function to return the pmap impl.
This is a helper function to return the pmap impl .
def _get_pmap_impl(f, devices, has_tpu):
  """This is a helper function to return the pmap impl.

  Args:
    f: a function that takes ndarrays and returns ndarrays.
    devices: a list of strings; the device list.
    has_tpu: boolean; whether `devices` contains TPU devices.

  Returns:
    A function that takes tensors and returns tensors.
  """
  if has_tpu:
    # Workaround b/121383831
    output_is_list = [False]  # Use list for mutability
    def recorder(args, kwargs, res):
      del args, kwargs
      output_is_list[0] = isinstance(res, list)
      return res
    f = _record_result_type(recorder, f)

  def tf_f(*tf_args):
    """A wrapper for `f` that takes/returns tensors."""
    np_args = _tf_to_np(tf_args)
    np_out = f(*np_args)
    return np_out

  if has_tpu:

    @tf.function(autograph=False)
    def fn(inputs):
      # TODO(wangpeng): Supply the `device_assignment` argument to
      # tpu.replicate, calculated from `devices`.
      res = tf.compat.v1.tpu.replicate(tf_f, inputs)
      # Workaround b/121383831
      if (res and isinstance(res[0], list) and len(res[0]) == 1 and
          not output_is_list[0]):
        res = [x[0] for x in res]
      return res

    return fn
  else:
    # This is run in a tf.function so that the various underlying functions can
    # be run in parallel.
    # The trace happens on the client, so any devices should not depend on any
    # side effects.

    jit_tf_f = tf.function(tf_f, autograph=False)

    @tf.function(autograph=False)
    def fn(all_per_device_args):
      """Multi-device function with calls placed on the correct device."""

      results = []
      for per_device_args, device in zip(all_per_device_args, devices):
        with tf.device(device):
          results.append(jit_tf_f(*per_device_args))
      return results

    return fn
def _get_pmap_impl ( f , devices , has_tpu ) : if has_tpu : # Workaround b/121383831 output_is_list = [ False ] # Use list for mutability def recorder ( args , kwargs , res ) : del args , kwargs output_is_list [ 0 ] = isinstance ( res , list ) return res f = _record_result_type ( recorder , f ) def tf_f ( * tf_args ) : """A wrapper for `f` that takes/returns tensors.""" np_args = _tf_to_np ( tf_args ) np_out = f ( * np_args ) return np_out if has_tpu : @ tf . function ( autograph = False ) def fn ( inputs ) : # TODO(wangpeng): Supply the `device_assignment` argument to # tpu.replicate, calculated from `devices`. res = tf . compat . v1 . tpu . replicate ( tf_f , inputs ) # Workaround b/121383831 if ( res and isinstance ( res [ 0 ] , list ) and len ( res [ 0 ] ) == 1 and not output_is_list [ 0 ] ) : res = [ x [ 0 ] for x in res ] return res return fn else : # This is run in a tf.function so that the various underlying functions can # be run in parallel. # The trace happens on the client, so any devices should not depend on any # side effects. jit_tf_f = tf . function ( tf_f , autograph = False ) @ tf . function ( autograph = False ) def fn ( all_per_device_args ) : """Multi-device function with calls placed on the correct device.""" results = [ ] for per_device_args , device in zip ( all_per_device_args , devices ) : with tf . device ( device ) : results . append ( jit_tf_f ( * per_device_args ) ) return results return fn