(forward_and_backward_fn,
                          optimizer,
                          n_devices,
                          accelerate=True,
                          adasum=False)
return multi_device_update_fn
Accelerates the given forward_and_backward_fn function.
Accelerates the given forward_and_backward_fn function.
Accelerates the given forward_and_backward_fn function .
def _accelerate_update_fn(forward_and_backward_fn,
                          optimizer,
                          n_devices,
                          accelerate=True,
                          adasum=False):
  """Accelerates the given forward_and_backward_fn function."""
  if n_devices == 1:
    def single_device_update_fn(
        weights_and_slots, step, opt_params, batch, state, rng):
      step = jnp.array(step, dtype=jnp.int32)  # Needed in TFNP backend.
      weights, slots = weights_and_slots
      (loss, state), gradients = forward_and_backward_fn(
          batch, weights, state, rng)
      weights, slots, stats = optimizer.tree_update(
          step, gradients, weights, slots, opt_params, store_slots=False)
      stats['loss'] = loss
      return (weights, slots), state, stats
    if accelerate:
      # TODO(afrozm): Find out the status of buffer donation on GPUs, then do
      #  donate_argnums=(0,).
      single_device_update_fn = fastmath.jit(single_device_update_fn)
    return single_device_update_fn

  # More than one device (core), i.e. all of TPU configurations etc.
  assert n_devices > 1, f'{n_devices} should be greater than 1.'

  @functools.partial(fastmath.pmap, axis_name='batch', donate_argnums=(0,))
  def _multi_device_update_fn(
      weights_and_slots, step, opt_params, batch, state, rng):
    # All tensors should have the first dimension = n_devices.
    weights, slots = weights_and_slots
    (loss, state), gradients = (
        forward_and_backward_fn(batch, weights, state, rng))
    gradients = _average_multidevice_gradients(gradients, adasum=adasum)
    weights, slots, stats = optimizer.tree_update(
        step, gradients, weights, slots, opt_params, store_slots=False)
    stats['loss'] = loss
    return (weights, slots), state, stats

  def multi_device_update_fn(
      weights_and_slots, step, opt_params, batch, state, rng):
    # Need to replicate step to n_devices leading dimension.
    return _multi_device_update_fn(weights_and_slots,
                                   jnp.repeat(step, n_devices), opt_params,
                                   batch, state, rng)

  return multi_device_update_fn
def _accelerate_update_fn ( forward_and_backward_fn , optimizer , n_devices , accelerate = True , adasum = False ) : if n_devices == 1 : def single_device_update_fn ( weights_and_slots , step , opt_params , batch , state , rng ) : step = jnp . array ( step , dtype = jnp . int32 ) # Needed in TFNP backend. weights , slots = weights_and_slots ( loss , state ) , gradients = forward_and_backward_fn ( batch , weights , state , rng ) weights , slots , stats = optimizer . tree_update ( step , gradients , weights , slots , opt_params , store_slots = False ) stats [ 'loss' ] = loss return ( weights , slots ) , state , stats if accelerate : # TODO(afrozm): Find out the status of buffer donation on GPUs, then do #  donate_argnums=(0,). single_device_update_fn = fastmath . jit ( single_device_update_fn ) return single_device_update_fn # More than one device (core), i.e. all of TPU configurations etc. assert n_devices > 1 , f'{n_devices} should be greater than 1.' @ functools . partial ( fastmath . pmap , axis_name = 'batch' , donate_argnums = ( 0 , ) ) def _multi_device_update_fn ( weights_and_slots , step , opt_params , batch , state , rng ) : # All tensors should have the first dimension = n_devices. weights , slots = weights_and_slots ( loss , state ) , gradients = ( forward_and_backward_fn ( batch , weights , state , rng ) ) gradients = _average_multidevice_gradients ( gradients , adasum = adasum ) weights , slots , stats = optimizer . tree_update ( step , gradients , weights , slots , opt_params , store_slots = False ) stats [ 'loss' ] = loss return ( weights , slots ) , state , stats def multi_device_update_fn ( weights_and_slots , step , opt_params , batch , state , rng ) : # Need to replicate step to n_devices leading dimension. return _multi_device_update_fn ( weights_and_slots , jnp . repeat ( step , n_devices ) , opt_params , batch , state , rng ) return multi_device_update_fn