Attend the Women in ML Symposium on December 7 Register now

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.

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.