(layer, optimizer, n_devices,
                            stats_name=None, adasum=False)
return fbo
Create the fbo function for a given layer and optimizer.
Create the fbo function for a given layer and optimizer.
Create the fbo function for a given layer and optimizer .
def _fbo_with_layer_and_opt(layer, optimizer, n_devices,
                            stats_name=None, adasum=False):
  """Create the fbo function for a given layer and optimizer."""
  def fbo(inputs, weights, grads, state, slots, opt_params, rng, step):
    """FBO of the layer."""
    # We need a layer pure_fn but only for inputs and weights.
    def pure_fn_without_state_and_rng(x, w):
      return layer.pure_fn(x, w, state, rng)

    # Calculate the vector-Jacobian product of the reduced pure fn.
    activations, vjp_fn, new_state = fastmath.vjp(
        pure_fn_without_state_and_rng, inputs, weights, has_aux=True)

    # In the loss layer, set gradients to 1 with the dtype of activations=loss.
    if grads is None and stats_name is not None:
      grads = jnp.ones((), dtype=activations.dtype)

    # The vjp function returns gradients with respect to inputs and weights.
    grads_inputs, grads_weights = vjp_fn(grads)

    # For non-trainable layers, return the calculated arguments.
    if _is_empty_tuple(weights):
      stats = {}
      if stats_name is not None:
        stats[stats_name] = activations
      return weights, new_state, slots, grads_inputs, stats

    # In multi-device setting, average gradients from multiple devices.
    if n_devices > 1:
      grads_weights = _average_multidevice_gradients(
          grads_weights, adasum=adasum)

    # Run the optimizer.
    new_weights, new_slots, stats = optimizer.tree_update(
        step, grads_weights, weights, slots, opt_params, store_slots=False)
    if stats_name is not None:
      stats[stats_name] = activations
    return new_weights, new_state, new_slots, grads_inputs, stats

  return fbo
def _fbo_with_layer_and_opt ( layer , optimizer , n_devices , stats_name = None , adasum = False ) : def fbo ( inputs , weights , grads , state , slots , opt_params , rng , step ) : """FBO of the layer.""" # We need a layer pure_fn but only for inputs and weights. def pure_fn_without_state_and_rng ( x , w ) : return layer . pure_fn ( x , w , state , rng ) # Calculate the vector-Jacobian product of the reduced pure fn. activations , vjp_fn , new_state = fastmath . vjp ( pure_fn_without_state_and_rng , inputs , weights , has_aux = True ) # In the loss layer, set gradients to 1 with the dtype of activations=loss. if grads is None and stats_name is not None : grads = jnp . ones ( ( ) , dtype = activations . dtype ) # The vjp function returns gradients with respect to inputs and weights. grads_inputs , grads_weights = vjp_fn ( grads ) # For non-trainable layers, return the calculated arguments. if _is_empty_tuple ( weights ) : stats = { } if stats_name is not None : stats [ stats_name ] = activations return weights , new_state , slots , grads_inputs , stats # In multi-device setting, average gradients from multiple devices. if n_devices > 1 : grads_weights = _average_multidevice_gradients ( grads_weights , adasum = adasum ) # Run the optimizer. new_weights , new_slots , stats = optimizer . tree_update ( step , grads_weights , weights , slots , opt_params , store_slots = False ) if stats_name is not None : stats [ stats_name ] = activations return new_weights , new_state , new_slots , grads_inputs , stats return fbo