tff.framework.MergeableCompExecutionContext

Context which executes mergeable computations in subrounds.

Inherits From: SyncContext

tff.framework.MergeableCompExecutionContext(
    async_contexts: Sequence[tff.framework.AsyncContext],
    compiler_fn: Optional[Callable[[_Computation], MergeableCompForm]] = None,
    transform_args: Optional[Callable[[object], object]] = None,
    transform_result: Optional[Callable[[object], object]] = None,
    num_subrounds: Optional[int] = None
)

This context relies on retrying behavior of the underlying asynchronous execution contexts to localize retries to these subrounds. That is, if a subround fails, this subround is the only computation that is retried. This allows MergeableCompExecutionContext to execute larger rounds than a runtime which retries the entire round in the case of e.g. a worker failure.

async_contexts Sequence of TFF execution contexts. These contexts are assumed to implement their invoke method as a coroutine function, returning an awaitable.
compiler_fn An optional callable which accepts a tff.framework.ConcreteComputation and returns an instance of MergeableCompForm. If not provided, this context will only execute instances of MergeableCompForm directly.
transform_args An Optional Callable used to transform the args before they are passed to the computation.
transform_result An Optional Callable used to transform the result before it is returned.
num_subrounds An optional integer, specifying total the number of subrounds desired. If unspecified, the length of async_contexts will determine the number of subrounds. If more subrounds are requested than contexts are passed, invocations will be sequentialized.

Methods

invoke

View source

invoke(
    comp: Union[tff.framework.MergeableCompForm, tff.Computation],
    arg: Optional[object] = None
)

Invokes computation comp with argument arg.

Args
comp The computation being invoked. The Python type of comp expected here (e.g., pb.Computation. tff.framework.ConcreteComputation, or other) may depend on the context. It is the responsibility of the concrete implementation of this interface to verify that the type of comp matches what the context is expecting.
arg The argument passed to the computation. If no argument is passed, this will be None. Structural argument types will be normalized into tff.structure.Structs.

Returns
The result of invocation, which is context-dependent.