examples/research_projects/jax-projects/big_bird/evaluate.py/get_best_valid_start_end_idx
def get_best_valid_start_end_idx(start_scores, end_scores, top_k=1, max_size=100):
    best_start_scores, best_start_idx = jax.lax.top_k(start_scores, top_k)
    best_end_scores, best_end_idx = jax.lax.top_k(end_scores, top_k)

    widths = best_end_idx[:, None] - best_start_idx[None, :]
    mask = jnp.logical_or(widths < 0, widths > max_size)
    scores = (best_end_scores[:, None] + best_start_scores[None, :]) - (1e8 * mask)
    best_score = jnp.argmax(scores).item()

    return best_start_idx[best_score % top_k], best_end_idx[best_score // top_k]