p-doom/crowd-code-dataset-0.2 huggingface dataset:
[content] jasmine/train_dynamics.py (python)
def _calculate_top_k_accuracy(
    token_logits_BTNV: jax.Array,
    video_tokens_BTN: jax.Array,
    mask_BTN: jax.Array,
    k: int,
) -> jax.Array:
    _, topk_indices_BTNK = jax.lax.top_k(token_logits_BTNV, k)
    topk_correct = jnp.any(
        topk_indices_BTNK == video_tokens_BTN[..., jnp.newaxis], axis=-1
    )
    topk_acc = (mask_BTN * topk_correct).sum() / mask_BTN.sum()
    return topk_acc