def test_attr_type_proto(self): # type: () -> None
# type_proto
type = TypeProto()
attr = helper.make_attribute("type_proto", type)
self.assertEqual(attr.name, "type_proto")
self.assertEqual(attr.tp, type)
self.assertEqual(attr.type, AttributeProto.TYPE_PROTO)
# type_protos
types = [TypeProto(), TypeProto()]
attr = helper.make_attribute("type_protos", types)
self.assertEqual(attr.name, "type_protos")
self.assertEqual(list(attr.type_protos), types)
self.assertEqual(attr.type, AttributeProto.TYPE_PROTOS)
def make_sequence_type_proto(
inner_type_proto: TypeProto,
) -> TypeProto:
"""Makes a sequence TypeProto."""
type_proto = TypeProto()
type_proto.sequence_type.elem_type.CopyFrom(inner_type_proto)
return type_proto
def make_optional_type_proto(
inner_type_proto: TypeProto,
) -> TypeProto:
"""Makes an optional TypeProto."""
type_proto = TypeProto()
type_proto.optional_type.elem_type.CopyFrom(inner_type_proto)
return type_proto
def make_sparse_tensor_type_proto(
elem_type: int,
shape: Optional[Sequence[Union[Text, int, None]]],
shape_denotation: Optional[List[Text]] = None,
) -> TypeProto:
"""Makes a SparseTensor TypeProto based on the data type and shape."""
type_proto = TypeProto()
sparse_tensor_type_proto = type_proto.sparse_tensor_type
sparse_tensor_type_proto.elem_type = elem_type
sparse_tensor_shape_proto = sparse_tensor_type_proto.shape
if shape is not None:
# You might think this is a no-op (extending a normal Python
# list by [] certainly is), but protobuf lists work a little
# differently; if a field is never set, it is omitted from the
# resulting protobuf; a list that is explicitly set to be
# empty will get an (empty) entry in the protobuf. This
# difference is visible to our consumers, so make sure we emit
# an empty shape!
sparse_tensor_shape_proto.dim.extend([])
if shape_denotation:
if len(shape_denotation) != len(shape):
raise ValueError(
'Invalid shape_denotation. '
'Must be of the same length as shape.')
for i, d in enumerate(shape):
dim = sparse_tensor_shape_proto.dim.add()
if d is None:
pass
elif isinstance(d, int):
dim.dim_value = d
elif isinstance(d, str):
dim.dim_param = d
else:
raise ValueError(
'Invalid item in shape: {}. '
'Needs to be of int or text.'.format(d))
if shape_denotation:
dim.denotation = shape_denotation[i]
return type_proto