p-doom/crowd-code-dataset-0.2 huggingface dataset:
[content] train_tokenizer.py (python)

    def tokenizer_loss_fn(params, state, inputs):
        # --- Compute loss ---
        outputs = state.apply_fn(
            params, inputs, training=True, rngs={"dropout": inputs["rng"]}
        )
        mse = jnp.square(inputs["videos"] - outputs["recon"]).mean()
        q_loss = jnp.square(jax.lax.stop_gradient(outputs["emb"]) - outputs["z"]).mean()
        commitment_loss = jnp.square(
            outputs["emb"] - jax.lax.stop_gradient(outputs["z"])
        ).mean()
        loss = mse + q_loss + args.vq_beta * commitment_loss

        # --- Compute validation metrics ---
        gt = inputs["videos"].clip(0, 1).reshape(-1, *inputs["videos"].shape[2:])
        recon = outputs["recon"].clip(0, 1).reshape(-1, *outputs["recon"].shape[2:])
        psnr = pix.psnr(gt, recon).mean()
        ssim = pix.ssim(gt, recon).mean()
        _, index_counts = jnp.unique_counts(
            jnp.ravel(outputs["indices"]), size=args.num_latents, fill_value=0
        )
        codebook_usage = (index_counts != 0).mean()
        metrics = dict(
            loss=loss,
            mse=mse,
            q_loss=q_loss,
            commitment_loss=commitment_loss,
            psnr=psnr,
            ssim=ssim,
            codebook_usage=codebook_usage,
        )
        return loss, (outputs["recon"], metrics)

    @jax.jit
    def train_step(state, inputs):
        grad_fn = jax.value_and_grad(tokenizer_loss_fn, has_aux=True, allow_int=True)
        (loss, (recon, metrics)), grads = grad_fn(state.params, state, inputs)
        state = state.apply_gradients(grads=grads)
        if args.log_gradients:
            metrics["encoder_gradients_std/"] = jax.tree.map(
                lambda x: x.std(), grads["params"]["encoder"]
            )
            metrics["vq_gradients_std/"] = jax.tree.map(
                lambda x: x.std(), grads["params"]["vq"]
            )
            metrics["decoder_gradients_std/"] = jax.tree.map(
                lambda x: x.std(), grads["params"]["decoder"]
            )
        return state, loss, recon, metrics

    return tokenizer_loss_fn, train_step


if __name__ == "__main__":
    args = tyro.cli(Args)
[content] vscode.git.Git (log)
2025-06-25 00:38:20.780 [info] > git add -A -- /lustre/groups/haicu/workspace/franz.srambical/jafar/train_tokenizer.py [32ms]
2025-06-25 00:38:20.802 [info] > git config --get commit.template [10ms]
2025-06-25 00:38:20.817 [info] > git for-each-ref --format=%(refname)%00%(upstream:short)%00%(objectname)%00%(upstream:track)%00%(upstream:remotename)%00%(upstream:remoteref) refs/heads/seeded-episode-sampling refs/remotes/seeded-episode-sampling [16ms]
2025-06-25 00:38:20.848 [info] > git status -z -uall [19ms]
2025-06-25 00:38:20.859 [info] > git for-each-ref --sort -committerdate --format %(refname) %(objectname) %(*objectname) [21ms]
[content] vscode.git.Git (log)
2025-06-25 00:38:22.320 [info] > git ls-files --stage -- /lustre/groups/haicu/workspace/franz.srambical/jafar/train_tokenizer.py [10ms]
2025-06-25 00:38:22.331 [info] > git ls-tree -l HEAD -- /lustre/groups/haicu/workspace/franz.srambical/jafar/train_tokenizer.py [11ms]
2025-06-25 00:38:22.341 [info] > git cat-file -s e0b223487cd0fd5e6b69aae9a58441a50592c4db [10ms]
2025-06-25 00:38:22.379 [info] > git config --get commit.template [10ms]
2025-06-25 00:38:22.392 [info] > git for-each-ref --format=%(refname)%00%(upstream:short)%00%(objectname)%00%(upstream:track)%00%(upstream:remotename)%00%(upstream:remoteref) refs/heads/seeded-episode-sampling refs/remotes/seeded-episode-sampling [14ms]
2025-06-25 00:38:22.419 [info] > git status -z -uall [17ms]
2025-06-25 00:38:22.434 [info] > git for-each-ref --sort -committerdate --format %(refname) %(objectname) %(*objectname) [23ms]
[content] vscode.git.Git (log)
2025-06-25 00:38:23.538 [info] > git -c user.useConfigOnly=true commit --quiet --allow-empty-message --file - [739ms]
2025-06-25 00:38:23.538 [info] [WARNING] Unstaged files detected.
[INFO] Stashing unstaged files to /ictstr01/home/aih/franz.srambical/.cache/pre-commit/patch1750804703-3444332.
black....................................................................Passed
[INFO] Restored changes from /ictstr01/home/aih/franz.srambical/.cache/pre-commit/patch1750804703-3444332.
2025-06-25 00:38:23.555 [info] > git config --get commit.template [6ms]
2025-06-25 00:38:23.576 [info] > git config --get commit.template [10ms]
2025-06-25 00:38:23.592 [info] > git for-each-ref --format=%(refname)%00%(upstream:short)%00%(objectname)%00%(upstream:track)%00%(upstream:remotename)%00%(upstream:remoteref) refs/heads/seeded-episode-sampling refs/remotes/seeded-episode-sampling [17ms]
2025-06-25 00:38:23.622 [info] > git status -z -uall [18ms]
2025-06-25 00:38:23.637 [info] > git for-each-ref --sort -committerdate --format %(refname) %(objectname) %(*objectname) [24ms]
2025-06-25 00:38:23.664 [info] > git merge-base refs/heads/seeded-episode-sampling refs/remotes/origin/seeded-episode-sampling [12ms]
2025-06-25 00:38:23.684 [info] > git diff --name-status -z --diff-filter=ADMR 9e8a3860d8270b2dbf00311f48a6f8bf621b1aca...refs/remotes/origin/seeded-episode-sampling [12ms]
[content] vscode.git.Git (log)
2025-06-25 00:38:23.889 [info] > git config --local branch.seeded-episode-sampling.github-pr-owner-number [11ms]
2025-06-25 00:38:23.889 [warning] [Git][config] git config failed: Failed to execute git
2025-06-25 00:38:23.890 [info] > git diff --no-color [20ms]
[content] vscode.git.Git (log)
2025-06-25 00:38:23.900 [info] > git config --get commit.template [11ms]
2025-06-25 00:38:23.913 [info] > git for-each-ref --format=%(refname)%00%(upstream:short)%00%(objectname)%00%(upstream:track)%00%(upstream:remotename)%00%(upstream:remoteref) refs/heads/seeded-episode-sampling refs/remotes/seeded-episode-sampling [14ms]
2025-06-25 00:38:23.929 [info] > git log --format=%H%n%aN%n%aE%n%at%n%ct%n%P%n%D%n%B -z --shortstat --diff-merges=first-parent -n50 --skip=0 --topo-order --decorate=full --stdin [67ms]
2025-06-25 00:38:23.931 [info] > git status -z -uall [12ms]
2025-06-25 00:38:23.950 [info] > git for-each-ref --sort -committerdate --format %(refname) %(objectname) %(*objectname) [23ms]
[content] vscode.git.Git (log)
2025-06-25 00:38:25.145 [info] > git ls-tree -l HEAD -- /lustre/groups/haicu/workspace/franz.srambical/jafar/train_tokenizer.py [11ms]
2025-06-25 00:38:25.153 [info] > git ls-tree -l HEAD -- /lustre/groups/haicu/workspace/franz.srambical/jafar/train_tokenizer.py [11ms]
[content] vscode.git.Git (log)
2025-06-25 00:39:01.714 [info] > git config --get commit.template [10ms]
2025-06-25 00:39:01.729 [info] > git for-each-ref --format=%(refname)%00%(upstream:short)%00%(objectname)%00%(upstream:track)%00%(upstream:remotename)%00%(upstream:remoteref) refs/heads/seeded-episode-sampling refs/remotes/seeded-episode-sampling [16ms]
2025-06-25 00:39:01.757 [info] > git status -z -uall [17ms]
2025-06-25 00:39:01.773 [info] > git for-each-ref --sort -committerdate --format %(refname) %(objectname) %(*objectname) [22ms]
[content] vscode.git.Git (log)
2025-06-25 00:39:02.353 [info] > git add -A -- /lustre/groups/haicu/workspace/franz.srambical/jafar/utils/nn.py [15ms]
2025-06-25 00:39:02.373 [info] > git config --get commit.template [9ms]
2025-06-25 00:39:02.390 [info] > git for-each-ref --format=%(refname)%00%(upstream:short)%00%(objectname)%00%(upstream:track)%00%(upstream:remotename)%00%(upstream:remoteref) refs/heads/seeded-episode-sampling refs/remotes/seeded-episode-sampling [17ms]
2025-06-25 00:39:02.417 [info] > git status -z -uall [18ms]
2025-06-25 00:39:02.431 [info] > git for-each-ref --sort -committerdate --format %(refname) %(objectname) %(*objectname) [24ms]
[content] vscode.git.Git (log)
2025-06-25 00:39:03.936 [info] > git ls-files --stage -- /lustre/groups/haicu/workspace/franz.srambical/jafar/utils/nn.py [13ms]
2025-06-25 00:39:03.946 [info] > git ls-tree -l HEAD -- /lustre/groups/haicu/workspace/franz.srambical/jafar/utils/nn.py [15ms]
2025-06-25 00:39:03.955 [info] > git cat-file -s 2b106612ba0c4615ee46876eff4e52dcf8d121aa [10ms]
[content] vscode.git.Git (log)
2025-06-25 00:39:04.912 [info] > git -c user.useConfigOnly=true commit --quiet --allow-empty-message --file - [767ms]
2025-06-25 00:39:04.912 [info] [WARNING] Unstaged files detected.
[INFO] Stashing unstaged files to /ictstr01/home/aih/franz.srambical/.cache/pre-commit/patch1750804744-3445150.
black....................................................................Passed
[INFO] Restored changes from /ictstr01/home/aih/franz.srambical/.cache/pre-commit/patch1750804744-3445150.
2025-06-25 00:39:04.927 [info] > git config --get commit.template [7ms]
2025-06-25 00:39:04.941 [info] > git config --get commit.template [6ms]
2025-06-25 00:39:04.958 [info] > git for-each-ref --format=%(refname)%00%(upstream:short)%00%(objectname)%00%(upstream:track)%00%(upstream:remotename)%00%(upstream:remoteref) refs/heads/seeded-episode-sampling refs/remotes/seeded-episode-sampling [16ms]
2025-06-25 00:39:04.984 [info] > git status -z -uall [17ms]
2025-06-25 00:39:04.991 [info] > git for-each-ref --sort -committerdate --format %(refname) %(objectname) %(*objectname) [19ms]
2025-06-25 00:39:05.007 [info] > git merge-base refs/heads/seeded-episode-sampling refs/remotes/origin/seeded-episode-sampling [8ms]
2025-06-25 00:39:05.024 [info] > git diff --name-status -z --diff-filter=ADMR 9e8a3860d8270b2dbf00311f48a6f8bf621b1aca...refs/remotes/origin/seeded-episode-sampling [10ms]
[content] vscode.git.Git (log)
2025-06-25 00:39:05.347 [info] > git config --local branch.seeded-episode-sampling.github-pr-owner-number [12ms]
2025-06-25 00:39:05.347 [warning] [Git][config] git config failed: Failed to execute git
2025-06-25 00:39:05.348 [info] > git diff --no-color [21ms]
2025-06-25 00:39:05.356 [info] > git config --get commit.template [9ms]
2025-06-25 00:39:05.370 [info] > git for-each-ref --format=%(refname)%00%(upstream:short)%00%(objectname)%00%(upstream:track)%00%(upstream:remotename)%00%(upstream:remoteref) refs/heads/seeded-episode-sampling refs/remotes/seeded-episode-sampling [15ms]
2025-06-25 00:39:05.397 [info] > git log --format=%H%n%aN%n%aE%n%at%n%ct%n%P%n%D%n%B -z --shortstat --diff-merges=first-parent -n50 --skip=0 --topo-order --decorate=full --stdin [78ms]
2025-06-25 00:39:05.400 [info] > git status -z -uall [19ms]
2025-06-25 00:39:05.415 [info] > git for-each-ref --sort -committerdate --format %(refname) %(objectname) %(*objectname) [23ms]
[content] vscode.git.Git (log)
2025-06-25 00:39:06.641 [info] > git ls-tree -l HEAD -- /lustre/groups/haicu/workspace/franz.srambical/jafar/utils/nn.py [10ms]
2025-06-25 00:39:06.649 [info] > git ls-tree -l HEAD -- /lustre/groups/haicu/workspace/franz.srambical/jafar/utils/nn.py [9ms]
2025-06-25 00:39:06.797 [info] > git config --get commit.template [10ms]
2025-06-25 00:39:06.812 [info] > git for-each-ref --format=%(refname)%00%(upstream:short)%00%(objectname)%00%(upstream:track)%00%(upstream:remotename)%00%(upstream:remoteref) refs/heads/seeded-episode-sampling refs/remotes/seeded-episode-sampling [16ms]
2025-06-25 00:39:06.837 [info] > git status -z -uall [15ms]
[content] vscode.git.Git (log)
2025-06-25 00:39:06.848 [info] > git for-each-ref --sort -committerdate --format %(refname) %(objectname) %(*objectname) [20ms]
2025-06-25 00:39:06.861 [info] > git show --textconv HEAD:utils/nn.py [14ms]
2025-06-25 00:39:06.894 [info] > git show --textconv HEAD:utils/nn.py [13ms]
[content] vscode.git.Git (log)
2025-06-25 00:39:11.695 [info] > git push origin seeded-episode-sampling:seeded-episode-sampling [1503ms]
2025-06-25 00:39:11.695 [info] To github.com:p-doom/jafar.git
   9e8a386..917a739  seeded-episode-sampling -> seeded-episode-sampling
2025-06-25 00:39:11.719 [info] > git config --get commit.template [12ms]
2025-06-25 00:39:11.730 [info] > git for-each-ref --format=%(refname)%00%(upstream:short)%00%(objectname)%00%(upstream:track)%00%(upstream:remotename)%00%(upstream:remoteref) refs/heads/seeded-episode-sampling refs/remotes/seeded-episode-sampling [12ms]
2025-06-25 00:39:11.750 [info] > git status -z -uall [13ms]
2025-06-25 00:39:11.770 [info] > git for-each-ref --sort -committerdate --format %(refname) %(objectname) %(*objectname) [23ms]
2025-06-25 00:39:11.796 [info] > git merge-base refs/heads/seeded-episode-sampling refs/remotes/origin/seeded-episode-sampling [11ms]
2025-06-25 00:39:11.817 [info] > git diff --name-status -z --diff-filter=ADMR 917a7397f06657ab25fa0612b8c9f8e00aefbf33...refs/remotes/origin/seeded-episode-sampling [11ms]
2025-06-25 00:39:11.874 [info] > git config --get commit.template [12ms]
2025-06-25 00:39:11.882 [info] > git for-each-ref --format=%(refname)%00%(upstream:short)%00%(objectname)%00%(upstream:track)%00%(upstream:remotename)%00%(upstream:remoteref) refs/heads/seeded-episode-sampling refs/remotes/seeded-episode-sampling [11ms]
[content] vscode.git.Git (log)
2025-06-25 00:39:11.908 [info] > git status -z -uall [15ms]
2025-06-25 00:39:11.924 [info] > git for-each-ref --sort -committerdate --format %(refname) %(objectname) %(*objectname) [22ms]
2025-06-25 00:39:12.015 [info] > git config --local branch.seeded-episode-sampling.github-pr-owner-number [4ms]
2025-06-25 00:39:12.015 [warning] [Git][config] git config failed: Failed to execute git
2025-06-25 00:39:12.016 [info] > git diff --no-color [12ms]
2025-06-25 00:39:12.083 [info] > git log --format=%H%n%aN%n%aE%n%at%n%ct%n%P%n%D%n%B -z --shortstat --diff-merges=first-parent -n50 --skip=0 --topo-order --decorate=full --stdin [84ms]