def _reverse_and_fbo_with_layer_and_opt(layer, optimizer, n_devices, adasum):
  """Create the reverse_and_fbo function for a given layer and optimizer."""
  def reverse_and_fbo(output, weights, grads, state, new_state,
                      slots, opt_params, rng, step):
    """Reverse and FBO of the layer."""
    # Call the reverse_and_grad method of the layer.
    inputs, (grads_inputs, grads_weights) = layer.reverse_and_grad(
        output, grads, weights, state, new_state, rng=rng)

    # For non-trainable layers, return the calculated arguments.
    if _is_empty_tuple(weights):
      return weights, slots, inputs, grads_inputs, {}

    # 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)

    return new_weights, new_slots, inputs, grads_inputs, stats

  return reverse_and_fbo
Create the reverse_and_fbo function for a given layer and optimizer.
trax/optimizers/trainer.py