def chain(*args):
  """Applies a list of chainable update transformations.

  Given a sequence of chainable transforms, `chain` returns an `init_fn`
  that constructs a `state` by concatenating the states of the individual
  transforms, and returns an `update_fn` which chains the update transformations
  feeding the appropriate state to each.

  Args:
    *args: a sequence of chainable (init_fn, update_fn) tuples.

  Returns:
    A single (init_fn, update_fn) tuple.
  """

  init_fns, update_fns = zip(*args)

  def init(params):
    return [fn(params) for fn in init_fns]

  def update(updates, state):
    new_state = []
    for s, fn in zip(state, update_fns):
      updates, new_s = fn(updates, s)
      new_state.append(new_s)
    return updates, new_state

  return init, update