Decorates/wraps functions as TFF TensorFlow computations.
tff.experimental_tf_fn_computation(
*args, tff_internal_types=None, **kwargs
)
This symbol can be used as either a decorator or a wrapper applied to a function given to it as an argument.
The supported patterns and examples of usage are as follows:
Convert an existing function inline into a TFF computation. This is the simplest mode of usage, and how one can embed existing non-TFF code for use with the TFF framework. In this mode, one invokes
tff.experimental_tf_fn_computation
with a pair of arguments, the first being a function that contains the logic, and the second being the TFF type of the parameter:foo = tff.experimental_tf_fn_computation(lambda x: x > 10, tf.int32)
After executing the above code snippet,
foo
becomes an instance of the abstract base classComputation
. Like all computations, it has thetype_signature
property:str(foo.type_signature)
'(int32 -> bool)'
``` The function passed as a parameter doesn't have to be a lambda, it can be any Python callable. One notable exception is that TFF does not handle arguments with default values. If one intends to create a computation that doesn't accept any arguments, the type argument is simply omitted. The function must be a no-argument function as well: ```python foo = tff.experimental_tf_fn_computation(lambda: tf.constant(10)) str(foo.type_signature) ``` >>> '( -> tf.int32)'
Decorate a callable with a TFF type to wrap it as a TFF computation. The only difference between this mode of usage and the one mentioned above is that instead of passing the callable as an argument,
tff.experimetnal_tf_func_computation
along with the optional type specifier is written above the callable's body.Here's an example of a computation that accepts a parameter:
@tff.experimental_tf_fn_computation(tf.int32) def foo(x): return x > 10
One can think of this mode of usage as merely a syntactic sugar for the example already given earlier:
foo = tff.tf_computation(lambda x: x > 10, tf.int32)
Here's an example of a no-parameter computation:
@tff.tf_computation def foo(): return tf.constant(10)
Again, this is merely syntactic sugar for the example given earlier:
foo = tff.tf_computation(lambda: tf.constant(10))
If the Python callable has multiple decorators,
tff.tf_computation
should be the outermost decorator (the one that appears first, or at the top).Create a polymorphic callable to be instantiated based on arguments, similarly to
tf.function
s that have been defined without an input signature.This mode of usage is symmetric to those above. One simply omits the type specifier, and applies
tff.experimental_tf_fn_computation
as a decorator or wrapper to a function/defun that does expect parameters.Here's an example of wrapping a lambda as a polymorphic callable:
foo = tff.tf_computation(lambda x, y: x > y)
The resulting
foo
can be used in the same ways as if it were had the type been declared; the corresponding computation is simply created on demand, in the same way as how polymorphictf.function
s create and cache concrete function definitions for each combination of argument types....foo(1, 2)... ...foo(0.5, 0.3)...
Here's an example of creating a polymorphic callable via decorator:
@tff.tf_computation def foo(x, y): return x > y
The syntax is symmetric to all examples already shown.
Args | |
---|---|
*args
|
Either a python callable, or TFF type spec, or both (with callable first), or neither, as documented in the 3 patterns and examples of usage above. |
Returns | |
---|---|
If invoked with a function as an argument, returns an instance of a TFF computation constructed based on this function. If called without one, as in the typical decorator style of usage, returns a callable that expects to be called with the function definition supplied as a parameter; see the patterns and examples of usage above. |