tff.templates.IterativeProcess

A process that includes an initialization and iterated computation.

tff.templates.IterativeProcess(
    initialize_fn: tff.Computation,
    next_fn: tff.Computation,
    next_is_multi_arg: Optional[bool] = None
)

Used in the notebooks

Used in the tutorials

An iterated process will usually be driven by a control loop like:

def initialize_fn():
  ...

def next_fn(state):
  ...

iterative_process = IterativeProcess(initialize_fn, next_fn)
state = iterative_process.initialize()
for round in range(num_rounds):
  state = iterative_process.next(state)

The initialize_fn function must return an object which is expected as input to and returned by the next_fn function. By convention, we refer to this object as state.

The iteration step (next_fn function) can accept arguments in addition to state (which must be the first argument), and return additional arguments, with state being the first output argument:

def next_fn(state, round_num):
  ...

iterative_process = ...
state = iterative_process.initialize()
for round in range(num_rounds):
  state, output = iterative_process.next(state, round)

initialize_fn A no-arg tff.Computation that returns the initial state of the iterative process. Let the type of this state be called S.
next_fn A tff.Computation that represents the iterated function. The first or only argument must be a type that is assignable from the state type S (tff.types.Type.is_assignable_from must return True). The first or only return value must also be assignable to the first or only argument, the same requirement as the S type.
next_is_multi_arg An optional boolean indicating that next_fn will receive more than just the state argument (if True) or only the state argument (if False). This parameter is primarily used to provide better error messages.

TypeError If initialize_fn and next_fn are not instances of tff.Computation.
TemplateInitFnParamNotEmptyError If initialize_fn has any input arguments.
TemplateStateNotAssignableError If the state returned by either initialize_fn or next_fn is not assignable to the first input argument of next_fn.

initialize A no-arg tff.Computation that returns the initial state.
next A tff.Computation that produces the next state.

Its first argument should always be the current state (originally produced by tff.templates.IterativeProcess.initialize), and the first (or only) returned value is the updated state.

state_type The tff.Type of the state of the process.