p-doom/crowd-code-dataset-0.2 huggingface dataset:
[content] jasmine/train_dynamics.py (python)
def _calculate_step_metrics(
    outputs: dict[str, jax.Array],
    gt: jax.Array,
    num_actions: int,
    num_patch_latents: int,
) -> tuple[jax.Array, dict]:
    mask = outputs["mask"]
    outputs["token_logits"] = outputs["token_logits"].astype(jnp.float32)
    ce_loss = optax.softmax_cross_entropy_with_integer_labels(
        outputs["token_logits"], outputs["video_tokens"]
    )
    ce_loss = (mask * ce_loss).sum() / mask.sum()
    acc = outputs["token_logits"].argmax(-1) == outputs["video_tokens"]
    acc = (mask * acc).sum() / mask.sum()
    select_probs = jax.nn.softmax(outputs["token_logits"])
    gt_val = gt.clip(0, 1).reshape(-1, *gt.shape[2:])
    recon = outputs["recon"].clip(0, 1).reshape(-1, *outputs["recon"].shape[2:])
    psnr = jnp.asarray(pix.psnr(gt_val, recon)).mean()
    ssim = jnp.asarray(pix.ssim(gt_val, recon)).mean()
    _, index_counts_tokenizer = jnp.unique_counts(
        jnp.ravel(outputs["video_tokens"]),
        size=num_patch_latents,
        fill_value=0,
    )
    codebook_usage_tokenizer = (index_counts_tokenizer != 0).mean()
    metrics = dict(
        cross_entropy_loss=ce_loss,
        masked_token_accuracy=acc,
[content] jasmine/train_dynamics.py (python)
        logits = outputs["token_logits"].astype(jnp.float32)
        mask = outputs["mask"].astype(jnp.float32)
        z_val = jax.nn.logsumexp(logits, axis=-1)
        z_loss = (mask * (z_val ** 2)).sum() / mask.sum()
        total_loss = ce_loss + args.z_loss_weight * z_loss
        metrics["z_loss"] = z_loss
        metrics["total_loss"] = total_loss
        return total_loss, (outputs["recon"], metrics)