def printable_type(t: TypeProto) -> Text:
if t.WhichOneof('value') == "tensor_type":
s = TensorProto.DataType.Name(t.tensor_type.elem_type)
if t.tensor_type.HasField('shape'):
if len(t.tensor_type.shape.dim):
s += str(', ' + 'x'.join(map(printable_dim, t.tensor_type.shape.dim)))
else:
s += str(', scalar')
return s
if t.WhichOneof('value') is None:
return ""
return 'Unknown type {}'.format(t.WhichOneof('value'))
def _extract_shape(type_proto: TypeProto):
which_value = type_proto.WhichOneof('value')
if which_value == 'tensor_type':
tensor = type_proto.tensor_type
else:
raise ValueError
return [
dim.dim_value if dim.WhichOneof('value') == 'dim_value' else None
for dim in tensor.shape.dim[1:]
]
