//  [python]
def _stitch(features):
  """Stitch features on the first dimension."""
  full_mask = tf.greater(features['task'], 1)
  step_mask = tf.reduce_any(full_mask, axis=-1)
  step_mask_exclude_last = tf.pad(step_mask,
                                  [[0, 0], [0, 1]],
                                  constant_values=False)[:, 1:]
  num_sequences = common_layers.shape_list(features['task'])[0]
  num_steps = common_layers.shape_list(features['task'])[1]
  connectors = tf.constant(PADDED_CONCATENATORS)
  # Select connectors
  connector_indices = tf.random.uniform(
      [num_sequences * num_steps], minval=0,
      maxval=len(PADDED_CONCATENATORS), dtype=tf.int32)
  selected_connectors = tf.reshape(
      tf.gather(connectors, connector_indices),
      [num_sequences, num_steps, len(PADDED_CONCATENATORS[0])])
  selected_connectors = tf.multiply(
      selected_connectors,
      tf.expand_dims(tf.cast(step_mask_exclude_last, tf.int32), 2),
      name='connector_mask')
  features['task'] = tf.concat([features['task'], selected_connectors], axis=-1)
  ref_offsets = tf.expand_dims(
      tf.cumsum(tf.reduce_sum(tf.cast(tf.greater(features['task'], 1), tf.int32), -1),
                exclusive=True, axis=-1), 2)
  features['task'] = tf.reshape(features['task'], [num_sequences, -1])
  full_mask = tf.greater(features['task'], 1)
  full_mask_int = tf.cast(full_mask, tf.int32)
  indices = tf.where(tf.sequence_mask(lengths=tf.reduce_sum(full_mask_int, -1)))
  values = tf.boolean_mask(tf.reshape(features['task'], [-1]),
                           tf.reshape(full_mask, [-1]))
  sparse_task = tf.sparse.SparseTensor(
      indices=indices, values=values,
      dense_shape=tf.cast(tf.shape(features['task']), tf.int64))
  # Stitch task and raw_task
  stitched_features = {}
  stitched_features['task'] = tf.sparse_tensor_to_dense(sparse_task)
  max_len = tf.reduce_max(
      tf.reduce_sum(tf.cast(tf.greater(stitched_features['task'], 1), tf.int32), -1))
  stitched_features['task'] = stitched_features['task'][:, :max_len]
  if 'raw_task' in features:
    connector_strs = tf.reshape(
        tf.gather(tf.constant(CONCATENATORS_STR), connector_indices),
        [num_sequences, num_steps])
    masked_connector_strs = tf.where(
        step_mask_exclude_last,
        connector_strs, tf.fill(tf.shape(connector_strs), ''))
    stitched_features['raw_task'] = tf.strings.reduce_join(
        tf.strings.reduce_join(tf.concat([
            tf.expand_dims(features['raw_task'], 2),
            tf.expand_dims(masked_connector_strs, 2)], axis=2), axis=-1), -1)
  # Stitch screen sequences
  action_lengths = tf.reduce_sum(tf.cast(
      tf.greater(features['verb_refs'][:, :, 0, 1],
                 features['verb_refs'][:, :, 0, 0]), tf.int32), -1)
  max_action_length = tf.reduce_max(action_lengths)
  def _pad(tensor, padding_value=0):
    shape_list = common_layers.shape_list(tensor)
    assert len(shape_list) >= 2
    padding_list = [[0, 0], [0, 1]] + [[0, 0]] * (len(shape_list) - 2)
    return tf.pad(tensor[:, :max_action_length],
                  padding_list, constant_values=padding_value)
  for key in features.keys():
    if key.endswith('_refs'):
      features[key] = tf.squeeze(features[key], 2)
      ref_mask = tf.expand_dims(tf.cast(
          tf.not_equal(features[key][:, :, 0],
                       features[key][:, :, 1]), tf.int32), 2)
      stitched_features[key] = tf.multiply(
          (features[key] + ref_offsets), ref_mask, name='ref_mask')
      stitched_features[key] = _pad(stitched_features[key])
    elif key in ['verbs', 'objects', 'consumed', 'obj_dom_pos',
                 'obj_text', 'obj_type', 'obj_clickable', 'obj_screen_pos',
                 'verb_refs', 'obj_refs', 'input_refs', 'obj_dom_dist']:
      features[key] = tf.squeeze(features[key], 2)
      stitched_features[key] = features[key]
      stitched_features[key] = _pad(
          stitched_features[key],
          padding_value=-1 if key == 'obj_type' else 0)
    elif key not in ['task', 'raw_task']:
      stitched_features[key] = features[key][:, 0]
  # Append eos to 'task'
  stitched_features['task'] = tf.pad(stitched_features['task'],
                                     [[0, 0], [0, 1]])
  task_mask = tf.cast(tf.greater(stitched_features['task'], 1), tf.int32)
  task_eos_mask = tf.pad(task_mask, [[0, 0], [1, 0]], constant_values=1)[:, :-1]
  stitched_features['task'] = stitched_features['task'] + (
      task_eos_mask - task_mask)
  # Append eos
  verb_mask = tf.cast(tf.greater(stitched_features['verbs'], 1), tf.int32)
  verb_eos_mask = tf.pad(verb_mask, [[0, 0], [1, 0]], constant_values=1)[:, :-1]
  verb_eos = verb_eos_mask - verb_mask
  stitched_features['verbs'] = stitched_features['verbs'] + verb_eos
  # Append last step refs to 'verb_refs'
  task_lengths = tf.where(tf.equal(stitched_features['task'], 1))[:, 1]
  eos_pos = tf.cast(tf.stack([task_lengths, task_lengths + 1], axis=1), tf.int32)
  action_mask = tf.cast(
      tf.sequence_mask(action_lengths, max_action_length + 1), tf.int32)
  action_and_eos_mask = tf.pad(action_mask, [[0, 0], [1, 0]],
                               constant_values=1)[:, :-1]
  verb_ref_eos = action_and_eos_mask - action_mask
  eos_refs = tf.multiply(
      tf.tile(tf.expand_dims(eos_pos, 1), [1, max_action_length + 1, 1]),
      tf.expand_dims(verb_ref_eos, 2), name='verb_ref_eos')
  stitched_features['verb_refs'] += eos_refs
  return stitched_features