def __init__(self, value, scope=None, type_spec=None): """Creates an instance of a value embedded in a lambda executor. The internal representation of a value can take one of the following supported forms: * An instance of `executor_value_base.ExecutorValue` that represents a value embedded in the target executor (functional or non-functional). * An as-yet unprocessed instance of `pb.Computation` that represents a function yet to be invoked (always a value of a functional type; any non-functional constructs should be processed on the fly). * A coroutine callable in Python that accepts a single argument that must be an instance of `LambdaExecutorValue` (or `None`), and that returns a result that is also an instance of `LambdaExecutorValue`. The associated type signature is always functional. * A single-level tuple (`anonymous_tuple.AnonymousTuple`) of instances of this class (of any of the supported forms listed here). Args: value: The internal representation of a value, as specified above. scope: An optional scope for computations. Only allowed if `value` is an unprocessed instance of `pb.Computation`, otherwise it must be `None` (the scope is meaningless in other cases). type_spec: An optional type signature, only allowed if `value` is a callable that represents a function (in which case it must be an instance of `computation_types.FunctionType`), otherwise it must be `None` (the type is implied in other cases). """ if isinstance(value, executor_value_base.ExecutorValue): py_typecheck.check_none(scope) py_typecheck.check_none(type_spec) type_spec = value.type_signature # pytype: disable=attribute-error elif isinstance(value, pb.Computation): if scope is not None: py_typecheck.check_type(scope, LambdaExecutorScope) py_typecheck.check_none(type_spec) type_spec = type_utils.get_function_type( type_serialization.deserialize_type(value.type)) elif callable(value): py_typecheck.check_none(scope) py_typecheck.check_type(type_spec, computation_types.FunctionType) else: py_typecheck.check_type(value, anonymous_tuple.AnonymousTuple) py_typecheck.check_none(scope) py_typecheck.check_none(type_spec) type_elements = [] for k, v in anonymous_tuple.iter_elements(value): py_typecheck.check_type(v, LambdaExecutorValue) type_elements.append((k, v.type_signature)) type_spec = computation_types.NamedTupleType([ (k, v) if k is not None else v for k, v in type_elements ]) self._value = value self._scope = scope self._type_signature = type_spec
async def _evaluate(self, comp, scope=None): """Evaluates or partially evaluates `comp` in `scope`. Args: comp: An instance of `pb.Computation` to process. scope: An optional `LambdaExecutorScope` to process it in, or `None` if there's no surrounding scope (the computation is self-contained). Returns: An instance of `LambdaExecutorValue` that isn't unprocessed (i.e., the internal representation directly in it isn't simply `comp` or any other instance of`pb.Computation`). The result, however, does not have, and often won't be processed completely; it suffices for this function to make only partial progress. """ py_typecheck.check_type(comp, pb.Computation) if scope is not None: py_typecheck.check_type(scope, LambdaExecutorScope) which_computation = comp.WhichOneof('computation') if which_computation in ['tensorflow', 'intrinsic', 'data', 'placement']: return LambdaExecutorValue(await self._target_executor.create_value( comp, type_utils.get_function_type( type_serialization.deserialize_type(comp.type)))) elif which_computation == 'lambda': def _make_comp_fn(scope, result, name, type_spec): async def _comp_fn(arg): return await self._evaluate(result, LambdaExecutorScope({name: arg}, scope)) return LambdaExecutorValue(_comp_fn, type_spec=type_spec) comp_lambda = getattr(comp, 'lambda') type_spec = type_utils.get_function_type( type_serialization.deserialize_type(comp.type)) return _make_comp_fn(scope, comp_lambda.result, comp_lambda.parameter_name, type_spec) elif which_computation == 'reference': return scope.resolve_reference(comp.reference.name) elif which_computation == 'call': if comp.call.argument.WhichOneof('computation') is not None: arg = LambdaExecutorValue(comp.call.argument, scope=scope) else: arg = None return await self.create_call( LambdaExecutorValue(comp.call.function, scope=scope), arg=arg) elif which_computation == 'selection': which_selection = comp.selection.WhichOneof('selection') return await self.create_selection( await self.create_call( LambdaExecutorValue(comp.selection.source, scope=scope)), **{which_selection: getattr(comp.selection, which_selection)}) elif which_computation == 'tuple': names = [str(e.name) if e.name else None for e in comp.tuple.element] values = [] for e in comp.tuple.element: val = LambdaExecutorValue(e.value, scope=scope) if (isinstance(val.type_signature, computation_types.FunctionType) and val.type_signature.parameter is None): val = self.create_call(val) else: async def _async_identity(x): return x val = _async_identity(val) values.append(val) values = await asyncio.gather(*values) return await self.create_tuple( anonymous_tuple.AnonymousTuple(list(zip(names, values)))) elif which_computation == 'block': for loc in comp.block.local: value = LambdaExecutorValue(loc.value, scope) scope = LambdaExecutorScope({loc.name: value}, scope) return await self._evaluate(comp.block.result, scope) else: raise NotImplementedError( 'Unsupported computation type "{}".'.format(which_computation))