def _set_type(name, feature, dtype):
  """Set the type of a Feature proto."""
  if dtype == tf.int64:
    feature.type = schema_pb2.INT
  elif dtype == tf.float32:
    feature.type = schema_pb2.FLOAT
  elif dtype == tf.string:
    feature.type = schema_pb2.BYTES
  else:
    raise ValueError('Feature "{}" has invalid dtype {}'.format(name, dtype))
Set the type of a Feature proto.
tensorflow_transform/tf_metadata/schema_utils.py