def validate_and_record_structure(self, features, labels):
      """Validates and records the structure of `features` and `labels`."""
      # Extract structure.
      has_labels = labels is not None
      feature_names = _extract_key_names(features)
      label_names = _extract_key_names(labels)

      if not self._initialized:
        # Record structure.
        self._initialized = True
        if self._feature_dims is not None:
          feature_dims_names = _extract_key_names(self._feature_dims)
          if feature_dims_names != feature_names:
            raise ValueError(
                'TPUConfig.input_partition_dims[0] mismatched feature'
                ' keys. Expected {}, got {}'.format(feature_names,
                                                    feature_dims_names))

          label_dims_names = _extract_key_names(self._label_dims)
          if self._label_dims is not None and label_dims_names != label_names:
            raise ValueError(
                'TPUConfig.input_partition_dims[1] mismatched label'
                ' keys. Expected {}, got {}'.format(label_names,
                                                    label_dims_names))

          self._flattened_input_dims = self._flatten_input_dims(
              self._feature_dims, feature_dims_names, self._label_dims,
              label_dims_names, label_names, has_labels)
Validates and records the structure of `features` and `labels`.
src/tpu_estimator.py