def _feature_spec_from_batched_tensors(
    tensors: Mapping[str, common_types.TensorType],
    is_evaluation_complete: bool,
    forced_sparse_keys: Iterable[str],
) -> Dict[str, common_types.FeatureSpecType]:
  """Infer a feature spec from a dict of tensors.

  Args:
    tensors: A dict whose keys are strings and values are `Tensor`,
      `SparseTensor`, or `RaggedTensor`s.
    is_evaluation_complete: A boolean indicating whether all analyzers have been
      evaluated or not.
    forced_sparse_keys: Keys of sparse tensors which should not be treated as
      varlen.

  Returns:
    A feature spec inferred from the types and shapes of the tensors.

  Raises:
    ValueError: If the feature spec cannot be inferred.
    TypeError: If any of the values of `tensors` are not a `Tensor`,
        `SparseTensor`, or `RaggedTensor`.
  """
  feature_spec = {}
  for name, tensor in tensors.items():
    if tensor.dtype not in (tf.string, tf.int64, tf.float32):
      raise ValueError(
          'Feature {} ({}) had invalid dtype {} for feature spec'.format(
              name, tensor, tensor.dtype))
    if isinstance(tensor, tf.SparseTensor):
      shape = tensor.get_shape()
      if shape.ndims > 2 or name in forced_sparse_keys:
        feature_spec[name] = tf.io.SparseFeature(
            index_key=[
                SPARSE_INDICES_NAME_TEMPLATE.format(
                    tensor_name=name, index=idx)
                for idx in range(shape.ndims - 1)
            ],
            value_key=SPARSE_VALUES_NAME_TEMPLATE.format(tensor_name=name),
            dtype=tensor.dtype,
            size=shape[1:],
            already_sorted=True)
      else:
        feature_spec[name] = tf.io.VarLenFeature(tensor.dtype)
    elif isinstance(tensor, tf.Tensor):
      shape = tensor.get_shape()
      if shape.ndims in [None, 0]:
        raise ValueError(
            'Feature {} ({}) had invalid shape {} for FixedLenFeature: must '
            'have rank at least 1'.format(name, tensor, shape))
      if is_evaluation_complete and any(
          dim is None for dim in shape.as_list()[1:]):
        raise ValueError(
            'Feature {} ({}) had invalid shape {} for FixedLenFeature: apart '
            'from the batch dimension, all dimensions must have known size'
            .format(name, tensor, shape))
      feature_spec[name] = tf.io.FixedLenFeature(shape.as_list()[1:],
                                                 tensor.dtype)
    elif isinstance(tensor, tf.RaggedTensor):
      feature_spec[name] = (
          _ragged_feature_spec_from_batched_tensor(name, tensor))
    else:
      raise TypeError(
          'Expected a Tensor, SparseTensor, or RaggedTensor got {} of type {} '
          'for feature {}'.format(tensor, type(tensor), name))

  return feature_spec
Infer a feature spec from a dict of tensors.

  Args:
    tensors: A dict whose keys are strings and values are `Tensor`,
      `SparseTensor`, or `RaggedTensor`s.
    is_evaluation_complete: A boolean indicating whether all analyzers have been
      evaluated or not.
    forced_sparse_keys: Keys of sparse tensors which should not be treated as
      varlen.

  Returns:
    A feature spec inferred from the types and shapes of the tensors.

  Raises:
    ValueError: If the feature spec cannot be inferred.
    TypeError: If any of the values of `tensors` are not a `Tensor`,
        `SparseTensor`, or `RaggedTensor`.
  
tensorflow_transform/schema_inference.py