def combine_and_expand(fragments: List['ObsFragment'], length: int,
                         **expand_args):
    """Expands each element of `fragments` and returns encoder input lists."""

    all_token_ids = []
    all_type_ids = []  # type: List[List[int]]
    all_float_features = []  # type: List[List[float]]

    for fragment in fragments:
      (token_ids, type_ids, float_features) = fragment.expand(**expand_args)
      all_token_ids.extend(token_ids)

      if not all_type_ids:
        all_type_ids = type_ids
      else:
        for old_type_ids, new_type_ids in zip(all_type_ids, type_ids):
          old_type_ids.extend(new_type_ids)

      if not all_float_features:
        all_float_features = float_features
      else:
        for old_features, new_features in zip(all_float_features,
                                              float_features):
          old_features.extend(new_features)

    if len(all_token_ids) > length:  # Truncate to `length`.
      all_token_ids = all_token_ids[:length]
      all_type_ids = [type_ids[:length] for type_ids in all_type_ids]
      all_float_features = [
          features[:length] for features in all_float_features
      ]

    padding_length = length - len(all_token_ids)
    if padding_length > 0:
      int_padding = [0] * padding_length
      float_padding = [0.0] * padding_length
      all_token_ids.extend(int_padding)
      all_type_ids = [type_ids + int_padding for type_ids in all_type_ids]
      all_float_features = [
          features + float_padding for features in all_float_features
      ]

    return all_token_ids, all_type_ids, all_float_features
Expands each element of `fragments` and returns encoder input lists.
language/search_agents/muzero/utils.py