def _flatten_input_dims(self, feature_dims, feature_dims_names, label_dims,
                            label_dims_names, label_names, has_labels):
      """Flatten input dims with the same order as flattened input tensors."""
      flattened_input_dims = []
      if feature_dims_names:
        # We need a fixed ordering for matching the tensors in features.
        flattened_input_dims.extend(
            [feature_dims[name] for name in feature_dims_names])
      else:
        flattened_input_dims.append(feature_dims)

      if label_dims_names:
        # We need a fixed ordering for matching the tensors in labels.
        flattened_input_dims.extend(
            [label_dims[name] for name in label_dims_names])
      else:
        if label_names:
          num_tensors_in_label = len(label_names)
        else:
          num_tensors_in_label = int(has_labels)
        # Setting `None` in input_partition_dims[1] will apply `None` to
        # all the tensors in labels, regardless of internal structure.
        flattened_input_dims.extend([label_dims] * num_tensors_in_label)

      return flattened_input_dims
Flatten input dims with the same order as flattened input tensors.
src/tpu_estimator.py