def _get_static_padded_shapes(model_config):
  """Return static padded shapes for TPU training."""

  max_input_symbols = model_config.data_options.max_num_tokens
  # Add 2 because of start and end symbol.
  max_decode_length = model_config.data_options.max_decode_length + 2

  padded_shapes = {
      constants.SOURCE_WORDPIECES_KEY: [max_input_symbols],
      constants.COPIABLE_INPUT_KEY: [max_input_symbols],
      constants.SOURCE_LEN_KEY: [],
      constants.TARGET_ACTION_TYPES_KEY: [max_decode_length],
      constants.TARGET_ACTION_IDS_KEY: [max_decode_length],
      constants.TARGET_LEN_KEY: [],
      constants.WEIGHT_KEY: [],
      constants.OUTPUT_TYPE_KEY: []
  }

  if model_config.model_parameters.use_foreign_key_features:
    padded_shapes[constants.FOREIGN_KEY_KEY] = [max_input_symbols]

  if model_config.model_parameters.use_segment_ids:
    padded_shapes[constants.SEGMENT_ID_KEY] = [max_input_symbols]

  if model_config.model_parameters.use_alignment_features:
    padded_shapes[constants.ALIGNED_KEY] = [max_input_symbols]

  return padded_shapes
Return static padded shapes for TPU training.
language/xsp/model/input_pipeline.py