Convert a `TensorInfo` object into a feature proto object.

def to_serialized_field(tensor_info):
  """Convert a `TensorInfo` object into a feature proto object."""
  # Select the type
  dtype = tensor_info.dtype

  # TODO(b/119937875): TF Examples proto only support int64, float32 and string
  # This create limitation like float64 downsampled to float32, bool converted
  # to int64 which is space ineficient, no support for complexes or quantized
  if tensor_info.dtype.is_integer or tensor_info.dtype.is_bool:
    dtype = tf.int64
  elif tensor_info.dtype.is_floating:
    dtype = tf.float32
  # It seems quite space inefficient to convert bool to int64
  # We may want to add support for complex, quantize dtype in the future

  # TFRecord only support 3 types
  if dtype not in (tf.int64, tf.float32, tf.string):
    raise NotImplementedError(
        'Serialization not implemented for {}'.format(dtype))

  # Select the feature proto type in function of the unknown shape
  if (tensor_info.shape is not None and  # Shape is a sequence (None, ...)
      tensor_info.shape.count(None) == 1 and
      tensor_info.shape[0] is None):
    return tf.io.FixedLenSequenceFeature(
        shape=tensor_info.shape[1:],
        dtype=dtype,
        allow_missing=True,
    )
  # At least one dimension is undefined
  elif tensor_info.shape is None or None in tensor_info.shape:
    return tf.io.VarLenFeature(dtype=dtype)
  else:
    return tf.io.FixedLenFeature(
        shape=tensor_info.shape,
        dtype=dtype,
    )