//  [python]
def sample_seq2seq(news_config: GroverConfig, initial_context, eos_token, ignore_ids=None, p_for_topp=0.95,
                   do_topk=False, max_len=1025):
    """
    Sample multiple outputs for a model in a seq2seq way.

    :param news_config: Configuration used to construct the model
    :param initial_context: [batch_size, seq_length] that we'll start generating with.
                            Invalid entries are padded.
    :param eos_token: Stop generating if you see this (tf scalar)
    :param ignore_ids: NEVER GENERATE THESE [vocab_size]
    :return:
    """
    batch_size, ctxb_end = get_shape_list(initial_context, expected_rank=2)
    # This just says 'ignore the pad character'
    if ignore_ids is None:
        ignore_ids = tf.constant([x == 0 for x in range(news_config.vocab_size)], dtype=tf.bool)

    with tf.name_scope('sample_sequence'):
        # Not everything might be the same size so we need to get lens

        lens = tf.reduce_sum(tf.cast(tf.not_equal(initial_context, news_config.pad_token_id), dtype=tf.int32), axis=1)

        seq_is_valid = tf.greater(lens, 0)
        ctxb_start = tf.reduce_min(tf.where(seq_is_valid, lens, ctxb_end * tf.ones_like(lens)))

        initial_ctx_part_a = tf.identity(initial_context[:, :ctxb_start])
        initial_ctx_part_b = tf.identity(initial_context[:, ctxb_start:])

        # Initial call to get cache
        context_output = sample_step(tokens=initial_ctx_part_a, ignore_ids=ignore_ids, news_config=news_config,
                                     batch_size=batch_size, p_for_topp=p_for_topp, cache=None, do_topk=do_topk)

        def _append_new_tokens(current_ctx, new_tokens):
            """ At each step we add tokens. Sometimes those tokens conflict with what we already have.
                This function fixes that. It doesnt fix probabilities though!"""
            current_ctx_len = get_shape_list(current_ctx, expected_rank=2)[1]

            new_tokens = tf.cond(
                current_ctx_len < ctxb_end,
                true_fn=lambda: tf.where(
                    tf.equal(initial_ctx_part_b[:, current_ctx_len - ctxb_start], news_config.pad_token_id),
                    new_tokens,
                    initial_ctx_part_b[:, current_ctx_len - ctxb_start]),
                false_fn=lambda: new_tokens,
            )
            # import ipdb
            # ipdb.set_trace()
            # if current_ctx_len < ctxb_end:
            #     existing_tokens = initial_ctx_part_b[:,current_ctx_len-ctxb_start]
            #
            #     new_tokens=tf.where(tf.equal(existing_tokens, news_config.pad_token_id),
            #              new_tokens,
            #              existing_tokens)

            return tf.concat([current_ctx, new_tokens[:, None]], 1)

        ctx = _append_new_tokens(initial_ctx_part_a, context_output['new_tokens'])
        cache = context_output['new_cache']
        probs = tf.concat([context_output['prev_probs'],
                           tf.batch_gather(context_output['new_probs_all'], ctx[:, -1,None])], 1)

        def body(ctx, cache, probs):
            """ for whatever reason this didn't work when I ran it on more than one at once... ugh."""
            next_outputs = sample_step(ctx[:, -1][:, None], ignore_ids=ignore_ids, news_config=news_config,
                                       batch_size=batch_size, p_for_topp=p_for_topp, cache=cache,
                                       do_topk=do_topk)

            # Update everything. We might need to use the old tokens.
            new_cache = tf.concat([cache, next_outputs['new_cache']], axis=-2)
            new_ids = _append_new_tokens(ctx, next_outputs['new_tokens'])
            new_probs = tf.concat([probs, tf.batch_gather(next_outputs['new_probs_all'], new_ids[:, -1,None])], 1)
            return [new_ids, new_cache, new_probs]

        def cond(ctx, cache, probs):
            is_eos = tf.equal(ctx, eos_token)
            seq_is_eos = tf.math.logical_or(tf.reduce_any(is_eos, axis=1), tf.math.logical_not(seq_is_valid))

            return tf.math.logical_not(tf.reduce_all(seq_is_eos))

        tokens, cache, probs = tf.while_loop(
            cond=cond, body=body, maximum_iterations=max_len - get_shape_list(ctx)[1],
            loop_vars=[ctx, cache, probs],
            shape_invariants=[tf.TensorShape([batch_size, None]),
                              tf.TensorShape(
                                  [batch_size, news_config.num_hidden_layers, 2,
                                   news_config.num_attention_heads,
                                   None, news_config.hidden_size // news_config.num_attention_heads]),
                              tf.TensorShape([batch_size, None]),
                              ],
            back_prop=False,
        )
    return tokens, probs