Attend the Women in ML Symposium on December 7 Register now

tff.types.to_type

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

Converts the argument into an instance of tff.Type.

Used in the notebooks

Used in the tutorials

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.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)
  ...

spec Either an instance of tff.Type, or an argument convertible to tff.Type.

An instance of tff.Type corresponding to the given spec.