Attend the Women in ML Symposium on December 7 Register now

tff.experimental_tf_fn_computation

Stay organized with collections Save and categorize content based on your preferences.

Decorates/wraps functions as TFF TensorFlow computations.

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:

  1. 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 class Computation. Like all computations, it has the type_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)'
    
    
  2. 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).

  3. Create a polymorphic callable to be instantiated based on arguments, similarly to tf.functions 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 polymorphic tf.functions 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 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.

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.