//  [python]
def pair_examples(ds: tf.data.Dataset, max_encoder_len, max_decoder_len, pool_size=5,
                  decoder_masks=None, encoder_masks=None) -> tf.data.Dataset:
  """
  return: tf.data.Dataset that yeilds pairs of examples that can be packed together
  """
  empty = {k: tf.zeros([pool_size] + list(spec.shape), spec.dtype) for
           k, spec in ds.element_spec.items()}
  blank_example = {k: tf.zeros(list(spec.shape), spec.dtype) for
                   k, spec in ds.element_spec.items()}
  empty["meta"] = tf.zeros([pool_size, 3], dtype=tf.int32)

  if decoder_masks is None:
    decoder_masks = []
    for key, v in get_target_modalities().items():
      mask_key = f"targets/{key}/mask"
      if mask_key not in ds.element_spec:
        raise ValueError(f"All targets must have a mask to do packing, but missing {mask_key}")
      decoder_masks.append(f"targets/{key}/mask")

  if encoder_masks is None:
    encoder_masks = []
    for key, v in get_input_modalities().items():
      encoder_masks.append((f"inputs/{key}/mask", v.get_static_sequence_len()))

  def _build_pair(pool, batch):
    out_seq_len = tf.add_n([tf.reduce_sum(tf.cast(batch[x], tf.int32)) for x in decoder_masks])

    input_seq_lens = []
    for mask, seq_len in encoder_masks:
      mask = batch[mask]
      if seq_len:
        valid_frames = tf.reduce_any(tf.reshape(mask > 0, [mask.shape[0], -1]), -1)
        n_frames = tf.reduce_sum(tf.cast(valid_frames, tf.int32))
        input_seq_lens.append(n_frames * seq_len)
      else:
        assert len(mask.shape) == 1
        input_seq_lens.append(tf.reduce_sum(tf.cast(mask, tf.int32)))
    input_seq_len = tf.add_n(input_seq_lens)

    meta = pool["meta"]
    valid = meta[:, 2] == 1
    pool_lens = tf.reduce_sum(meta[:, :2], -1)
    paired_input_lens = input_seq_len + meta[:, 0]
    paired_target_lens = out_seq_len + meta[:, 1]
    can_be_packed = tf.logical_and(paired_input_lens <= max_encoder_len, paired_target_lens  <= max_decoder_len)
    can_be_packed = tf.logical_and(can_be_packed, valid)

    # tf.print("Enter pool",
    #          "in:", input_seq_len,
    #          "out:", out_seq_len,
    #          "lens", pool_lens, "meta", pool["meta"],
    #          "packable w/:", can_be_packed)
    if tf.logical_not(tf.reduce_any(can_be_packed)):
      # tf.print("Nothing can be packed")
      # Nothing in the pool can be paired with this example
      if tf.reduce_all(valid):
        # The pool is full
        # tf.print("The pool is full")
        if (out_seq_len + input_seq_len) >= tf.reduce_max(pool_lens):
          # Write batch, pool is unchanged
          # tf.print("Is longer than everything in the pool, writing", batch["ixs"])
          output_pair = (batch, blank_example)
          output_pool = pool
        else:
          # Write the longest example in the pool and replace with `batch`
          # tf.print("Writing the longest example in pool", ix)
          ix = tf.argmax(pool_lens)

          output_batch = {k: v[ix] for k, v in pool.items() if k not in {"meta"}}
          output_pool = _write_to_pool(batch, pool, input_seq_len, out_seq_len, ix)
          output_pair = (output_batch, blank_example)
      else:
        ix = tf.where(tf.logical_not(valid))[0][0]
        # Add to the pool and write nothing
        # tf.print("Writing to pool at:", ix,  "write nothing")
        output_pool = _write_to_pool(batch, pool, input_seq_len, out_seq_len, ix)
        output_pair = (blank_example, blank_example)

    else:
      # pair with longest example in the pool that matches
      # tf.print("Pairing batch with", ix)
      ix = tf.argmax(pool_lens * (2*tf.cast(can_be_packed, tf.int32) - 1))
      other = {k: v[ix] for k, v in pool.items() if k != "meta"}
      output_pool = dict(pool)
      output_pool["meta"] = _set_vector(pool["meta"], ix, tf.constant([0, 0, 0]))
      output_pair = (batch, other)

    p1, p2 = output_pair
    return output_pool, {k: tf.stack([p1[k], p2[k]], 0) for k in batch}

  # Convert dataset into packable pairs
  ds = ds.scan(empty, _build_pair)

  # Remove blank pairs
  ds = ds.filter(lambda batch: tf.reduce_any(
    tf.stack([tf.reduce_any(tf.cast(batch[x], tf.bool)) for x in decoder_masks])))
  return ds