tff.types.to_type

View source on GitHub

Converts the argument into an instance of tff.Type.

tff.types.to_type(
    obj: object
) -> tff.types.Type

Used in the notebooks

Examples of arguments convertible to tensor types:

tf.int32
(tf.int32, [10])
(tf.int32, [None])
np.int32

Examples of arguments convertible to flat named tuple types:

[tf.int32, tf.bool]
(tf.int32, tf.bool)
[('a', tf.int32), ('b', tf.bool)]
('a', tf.int32)
collections.OrderedDict([('a', tf.int32), ('b', tf.bool)])

Examples of arguments convertible to nested named tuple types:

(tf.int32, (tf.float32, tf.bool))
(tf.int32, (('x', tf.float32), tf.bool))
((tf.int32, [1]), (('x', (tf.float32, [2])), (tf.bool, [3])))

attr.s class instances can also be used to describe TFF types by populating the fields with the corresponding types:

@attr.s(auto_attribs=True)
class MyDataClass:
  int_scalar: tf.Tensor
  string_array: tf.Tensor

  @classmethod
  def tff_type(cls) -> tff.Type:
    return tff.types.to_type(cls(
      int_scalar=tf.int32,
      string_array=tf.TensorSpec(dtype=tf.string, shape=[3]),
    ))

@tff.tf_computation(MyDataClass.tff_type())
def work(my_data):
  assert isinstance(my_data, MyDataClass)
  ...