tff.check_returns_type

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

Checks that the decorated function returns values of the provided type.

tff.check_returns_type(
    *args
)

Used in the notebooks

Used in the tutorials

This decorator can be used to ensure that a TFF computation returns a value of the expected type. For example:

@tff.tf_computation(tf.int32, tf.int32)
@tff.check_returns_type(tf.int32)
def add(a, b):
  return a + b

It can also be applied to non-TFF (Python) functions to ensure that the values they return conform to the expected type.

Note that this assertion is run whenever the function is called. In the case of @tff.tf_computation and @tff.federated_computations, this means that the assertion will run when the computation is traced. To enable this, @tff.check_returns_type should be applied inside the tff.tf_computation:

# YES:
@tff.tf_computation(...)
@tff.check_returns_type(...)
...

# NO:
@tff.check_returns_type(...) # Don't put this before the line below
@tff.tf_computation(...)
...

*args Either a Python function, or TFF type spec, or both (function first).

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 also tff.tf_computation for an extended documentation.