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.

Args
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.