p-doom/crowd-code-dataset-0.2 huggingface dataset:
[content] train_dynamics.py (python)
    dataloader = (
        jax.make_array_from_process_local_data(videos_sharding, elem)
        for elem in grain_iterator
    )
    if jax.process_index() == 0:
        first_videos = next(dataloader)
        sample_inputs = dict(videos=first_videos)
        compiled = train_step.lower(optimizer, sample_inputs).compile()
        print_compiled_memory_stats(compiled.memory_analysis())
        print_compiled_cost_analysis(compiled.cost_analysis())
        # Do not skip the first batch during training
        dataloader = itertools.chain([first_videos], dataloader)
    print(f"Starting training from step {step}...")
    while step < args.num_steps:
        for videos in dataloader:
            # --- Train step ---
            rng, _rng_mask = jax.random.split(rng, 2)
            inputs = dict(videos=videos, mask_rng=_rng_mask)
            loss, recon, metrics = train_step(optimizer.model, optimizer, inputs)
            metrics["lr"] = lr_schedule(step)
            print(f"Step {step}, loss: {loss}")
            step += 1

            # --- Logging ---
            if args.log:
                if step % args.log_interval == 0 and jax.process_index() == 0:
                    wandb.log(
                        {
                            "loss": loss,
                            "step": step,
                            **metrics,
                        }
                    )
                if step % args.log_image_interval == 0:
                    gt_seq = inputs["videos"][0].astype(jnp.float32) / 255.0
                    recon_seq = recon[0].clip(0, 1)
                    comparison_seq = jnp.concatenate((gt_seq, recon_seq), axis=1)
                    comparison_seq = einops.rearrange(
                        comparison_seq * 255, "t h w c -> h (t w) c"
                    )
                    if jax.process_index() == 0:
                        log_images = dict(
                            image=wandb.Image(np.asarray(gt_seq[args.seq_len - 1])),
                            recon=wandb.Image(np.asarray(recon_seq[args.seq_len - 1])),
                            true_vs_recon=wandb.Image(
                                np.asarray(comparison_seq.astype(np.uint8))
                            ),
                        )
                        wandb.log(log_images)
            # --- Checkpointing ---
            if args.save_ckpt and step % args.log_checkpoint_interval == 0:
                optimizer_state = nnx.state(optimizer)
                checkpoint_manager.save(
                    step,
                    args=ocp.args.Composite(
                        model_state=ocp.args.PyTreeSave(optimizer_state),  # type: ignore
                        dataloader_state=grain.checkpoint.CheckpointSave(  # type: ignore
                            grain_iterator  # type: ignore
                        ),
                    ),
                )
                print(f"Saved checkpoint at step {step}")
            if step >= args.num_steps:
                break

    checkpoint_manager.close()
[content] train_dynamics.py (python)
    if jax.process_index() == 0:
        first_videos = next(dataloader)
        sample_inputs = dict(videos=first_videos)
        compiled = train_step.lower(optimizer, sample_inputs).compile()
        print_compiled_memory_stats(compiled.memory_analysis())
        print_compiled_cost_analysis(compiled.cost_analysis())
        # Do not skip the first batch during training
        dataloader = itertools.chain([first_videos], dataloader)
    print(f"Starting training from step {step}...")
    jax.config.update("jax_transfer_guard", "disallow")
    should_break = False
    while step < args.num_steps and not should_break:
        for _ in range(5):
            # if step == 2:
            #     jax.profiler.start_trace("tensorboard")
            # --- Train step ---
            rng, _rng_mask = jax.random.split(rng, 2)
            inputs = dict(mask_rng=_rng_mask)
            loss, recon, metrics = train_step(optimizer, inputs)
            with jax.default_device(jax.devices("cpu")[0]):
                metrics["lr"] = float(lr_schedule(step))
            print(f"Step {step}, loss: {loss}")
            step += 1

            # --- Logging ---
            if args.log:
                if step % args.log_interval == 0 and jax.process_index() == 0:
                    wandb.log(
                        {
                            "loss": loss,
                            "step": step,
                            **metrics,
                        }
                    )
                if step % args.log_image_interval == 0:
                    pass
            # --- Checkpointing ---
            if args.save_ckpt and step % args.log_checkpoint_interval == 0:
                optimizer_state = nnx.state(optimizer)
                checkpoint_manager.save(
                    step,
                    args=ocp.args.Composite(
                        model_state=ocp.args.PyTreeSave(optimizer_state),  # type: ignore
                        dataloader_state=grain.checkpoint.CheckpointSave(  # type: ignore
                            grain_iterator  # type: ignore
                        ),
                    ),
                )
                print(f"Saved checkpoint at step {step}")
            if step >= args.num_steps:
                break

    checkpoint_manager.close()
    # jax.profiler.stop_trace()
[content] train_lam.py (python)
    @nnx.jit
    def train_step(
        optimizer: nnx.Optimizer,
        inputs: dict,
        action_last_active: jax.Array,
        rng: jax.Array,
    ) -> tuple[jax.Array, jax.Array, jax.Array, dict]:
        def loss_fn(
            model: LatentActionModel,
        ) -> tuple[jax.Array, tuple[jax.Array, jax.Array, dict]]:
            return lam_loss_fn(model, inputs)

        # --- Update model ---
        (loss, (recon, idx_counts, metrics)), grads = nnx.value_and_grad(
            loss_fn, has_aux=True
        )(optimizer.model)
        optimizer.update(grads)

        # --- Reset inactive latent actions ---
        codebook = optimizer.model.vq.codebook
        num_codes = len(codebook)
        active_codes = idx_counts != 0.0
        action_last_active = jnp.where(active_codes, 0, action_last_active + 1)
        p_code = active_codes / active_codes.sum()
        reset_idxs = jax.random.choice(rng, num_codes, shape=(num_codes,), p=p_code)
        do_reset = action_last_active >= args.vq_reset_thresh
        new_codebook = jnp.where(
            jnp.expand_dims(do_reset, -1), codebook[reset_idxs], codebook.value
        )
        optimizer.model.vq.codebook.value = new_codebook
        action_last_active = jnp.where(do_reset, 0, action_last_active)
        return loss, recon, action_last_active, metrics

    # --- TRAIN LOOP ---
    dataloader = (
        jax.make_array_from_process_local_data(videos_sharding, elem)
        for elem in grain_iterator
    )
    if jax.process_index() == 0:
        first_videos = next(dataloader)
        sample_inputs = dict(videos=first_videos)
        compiled = train_step.lower(optimizer, sample_inputs).compile()
        print_compiled_memory_stats(compiled.memory_analysis())
        print_compiled_cost_analysis(compiled.cost_analysis())
        # Do not skip the first batch during training
        dataloader = itertools.chain([first_videos], dataloader)
    print(f"Starting training from step {step}...")
    action_last_active = jnp.zeros(args.num_latents, dtype=jnp.int32)
[content] train_lam.py (python)
    @nnx.jit(donate_argnums=0)
    def train_step(
        optimizer: nnx.Optimizer,
        inputs: dict,
        action_last_active: jax.Array,
        rng: jax.Array,
    ) -> tuple[jax.Array, jax.Array, jax.Array, dict]:
        def loss_fn(
            model: LatentActionModel,
        ) -> tuple[jax.Array, tuple[jax.Array, jax.Array, dict]]:
            return lam_loss_fn(model, inputs)

        # --- Update model ---
        (loss, (recon, idx_counts, metrics)), grads = nnx.value_and_grad(
            loss_fn, has_aux=True
        )(optimizer.model)
        optimizer.update(grads)

        # --- Reset inactive latent actions ---
        codebook = optimizer.model.vq.codebook
        num_codes = len(codebook)
        active_codes = idx_counts != 0.0
        action_last_active = jnp.where(active_codes, 0, action_last_active + 1)
        p_code = active_codes / active_codes.sum()
        reset_idxs = jax.random.choice(rng, num_codes, shape=(num_codes,), p=p_code)
        do_reset = action_last_active >= args.vq_reset_thresh
        new_codebook = jnp.where(
            jnp.expand_dims(do_reset, -1), codebook[reset_idxs], codebook.value
        )
        optimizer.model.vq.codebook.value = new_codebook
        action_last_active = jnp.where(do_reset, 0, action_last_active)
        return loss, recon, action_last_active, metrics

    # --- TRAIN LOOP ---
    dataloader = (
        jax.make_array_from_process_local_data(videos_sharding, elem)
        for elem in grain_iterator
    )
    if jax.process_index() == 0:
        first_videos = next(dataloader)
        sample_inputs = dict(videos=first_videos)
        compiled = train_step.lower(optimizer, sample_inputs).compile()
        print_compiled_memory_stats(compiled.memory_analysis())
        print_compiled_cost_analysis(compiled.cost_analysis())
        # Do not skip the first batch during training
        dataloader = itertools.chain([first_videos], dataloader)
    print(f"Starting training from step {step}...")
    action_last_active = jnp.zeros(args.num_latents, dtype=jnp.int32)
[content] train_tokenizer.py (python)
    jax.config.update("jax_transfer_guard", "disallow")
    while step < args.num_steps:
        for videos in dataloader:
            # --- Train step ---
            inputs = dict(videos=videos)
            loss, recon, metrics = train_step(optimizer, inputs)
            # metrics["lr"] = lr_schedule(step)
            with jax.default_device("cpu"):
                step_on_cpu = jnp.array(step)
                metrics["lr"] = float(lr_schedule(step_on_cpu))
[content] train_lam.py (python)
    # @nnx.jit
    def train_step(
        optimizer: nnx.Optimizer,
        inputs: dict,
        action_last_active: jax.Array,
        rng: jax.Array,
    ) -> tuple[jax.Array, jax.Array, jax.Array, dict]:
        def loss_fn(
            model: LatentActionModel,
        ) -> tuple[jax.Array, tuple[jax.Array, jax.Array, dict]]:
            return lam_loss_fn(model, inputs)

        # --- Update model ---
        (loss, (recon, idx_counts, metrics)), grads = nnx.value_and_grad(
            loss_fn, has_aux=True
        )(optimizer.model)
        optimizer.update(grads)

        # --- Reset inactive latent actions ---
        codebook = optimizer.model.vq.codebook
        num_codes = len(codebook)
        active_codes = idx_counts != 0.0
        action_last_active = jnp.where(active_codes, 0, action_last_active + 1)
        p_code = active_codes / active_codes.sum()
        reset_idxs = jax.random.choice(rng, num_codes, shape=(num_codes,), p=p_code)
        do_reset = action_last_active >= args.vq_reset_thresh
        new_codebook = jnp.where(
            jnp.expand_dims(do_reset, -1), codebook[reset_idxs], codebook.value
        )
        optimizer.model.vq.codebook.value = new_codebook
        action_last_active = jnp.where(do_reset, 0, action_last_active)
        return loss, recon, action_last_active, metrics

    # --- TRAIN LOOP ---
    dataloader = (
        jax.make_array_from_process_local_data(videos_sharding, elem)
        for elem in grain_iterator
    )
    action_last_active = jnp.zeros(args.num_latents, dtype=jnp.int32)
    if jax.process_index() == 0:
        first_videos = next(dataloader)
        sample_inputs = dict(videos=first_videos)
        compiled = train_step.lower(
            optimizer, sample_inputs, action_last_active, rng
        ).compile()
[content] train_dynamics.py (python)
    @nnx.jit
    def train_step(
        optimizer: nnx.Optimizer, inputs: dict
    ) -> tuple[jax.Array, jax.Array, dict]:
        def loss_fn(model: Genie) -> tuple[jax.Array, tuple[jax.Array, dict]]:
            return dynamics_loss_fn(model, inputs)

        (loss, (recon, metrics)), grads = nnx.value_and_grad(loss_fn, has_aux=True)(
            optimizer.model
        )
        optimizer.update(grads)
        if args.log_gradients:
            metrics["gradients_std/"] = jax.tree.map(
                lambda x: x.std(), grads["params"]["dynamics"]
            )
        return loss, recon, metrics

    # --- TRAIN LOOP ---
    dataloader = (
        jax.make_array_from_process_local_data(videos_sharding, elem)
        for elem in grain_iterator
    )
    if jax.process_index() == 0:
        first_videos = next(dataloader)
        sample_inputs = dict(videos=first_videos)
        compiled = train_step.lower(optimizer, sample_inputs).compile()
        print_compiled_memory_stats(compiled.memory_analysis())
        print_compiled_cost_analysis(compiled.cost_analysis())
        # Do not skip the first batch during training
        dataloader = itertools.chain([first_videos], dataloader)
    print(f"Starting training from step {step}...")
    while step < args.num_steps:
        for videos in dataloader:
            # --- Train step ---
            rng, _rng_mask = jax.random.split(rng, 2)
            inputs = dict(videos=videos, mask_rng=_rng_mask)
            loss, recon, metrics = train_step(optimizer.model, optimizer, inputs)
[content] train_dynamics.py (python)
    @nnx.jit(donate_argnums=0)
    def train_step(
        optimizer: nnx.Optimizer, inputs: dict
    ) -> tuple[jax.Array, jax.Array, dict]:
        def loss_fn(model: Genie) -> tuple[jax.Array, tuple[jax.Array, dict]]:
            return dynamics_loss_fn(model, inputs)

        (loss, (recon, metrics)), grads = nnx.value_and_grad(loss_fn, has_aux=True)(
            optimizer.model
        )
        optimizer.update(grads)
        if args.log_gradients:
            metrics["gradients_std/"] = jax.tree.map(
                lambda x: x.std(), grads["params"]["dynamics"]
            )
        return loss, recon, metrics

    # --- TRAIN LOOP ---
    dataloader = (
        jax.make_array_from_process_local_data(videos_sharding, elem)
        for elem in grain_iterator
    )
    if jax.process_index() == 0:
        first_videos = next(dataloader)
        sample_inputs = dict(videos=first_videos)
        compiled = train_step.lower(optimizer, sample_inputs).compile()
        print_compiled_memory_stats(compiled.memory_analysis())
        print_compiled_cost_analysis(compiled.cost_analysis())
        # Do not skip the first batch during training
        dataloader = itertools.chain([first_videos], dataloader)
    print(f"Starting training from step {step}...")
    while step < args.num_steps:
        for videos in dataloader:
            # --- Train step ---
            rng, _rng_mask = jax.random.split(rng, 2)
            inputs = dict(videos=videos, mask_rng=_rng_mask)
            loss, recon, metrics = train_step(optimizer.model, optimizer, inputs)
[content] /fast/home/franz.srambical/jafar/train_dynamics.py (python)
    def train_step(
        optimizer: nnx.Optimizer, inputs: dict
    ) -> tuple[jax.Array, jax.Array, dict]:
        def loss_fn(model: Genie) -> tuple[jax.Array, tuple[jax.Array, dict]]:
            return dynamics_loss_fn(model, inputs)

        (loss, (recon, metrics)), grads = nnx.value_and_grad(loss_fn, has_aux=True)(
            optimizer.model
        )
        optimizer.update(grads)
        if args.log_gradients:
            metrics["gradients_std/"] = jax.tree.map(
                lambda x: x.std(), grads["params"]["dynamics"]
            )
        return loss, recon, metrics

    # --- TRAIN LOOP ---
    dataloader = (
        jax.make_array_from_process_local_data(videos_sharding, elem)
        for elem in grain_iterator
    )
    if jax.process_index() == 0:
        first_videos = next(dataloader)
        sample_inputs = dict(videos=first_videos)
        compiled = train_step.lower(optimizer, sample_inputs).compile()
        print_compiled_memory_stats(compiled.memory_analysis())
        print_compiled_cost_analysis(compiled.cost_analysis())
        # Do not skip the first batch during training
        dataloader = itertools.chain([first_videos], dataloader)
    print(f"Starting training from step {step}...")
    while step < args.num_steps:
        for videos in dataloader:
            # --- Train step ---
            rng, _rng_mask = jax.random.split(rng, 2)
            inputs = dict(videos=videos, mask_rng=_rng_mask)
            loss, recon, metrics = train_step(optimizer.model, optimizer, inputs)
[content] /fast/home/franz.srambical/jafar/train_dynamics.py (python)
    @nnx.jit
    def train_step(
        optimizer: nnx.Optimizer, inputs: dict
    ) -> tuple[jax.Array, jax.Array, dict]:
        def loss_fn(model: Genie) -> tuple[jax.Array, tuple[jax.Array, dict]]:
            return dynamics_loss_fn(model, inputs)

        (loss, (recon, metrics)), grads = nnx.value_and_grad(loss_fn, has_aux=True)(
            optimizer.model
        )
        optimizer.update(grads)
        if args.log_gradients:
            metrics["gradients_std/"] = jax.tree.map(
                lambda x: x.std(), grads["params"]["dynamics"]
            )
        return loss, recon, metrics

    # --- TRAIN LOOP ---
    dataloader = (
        jax.make_array_from_process_local_data(videos_sharding, elem)
        for elem in grain_iterator
    )
    if jax.process_index() == 0:
        first_videos = next(dataloader)
        sample_inputs = dict(videos=first_videos)
        compiled = train_step.lower(optimizer, sample_inputs).compile()
        print_compiled_memory_stats(compiled.memory_analysis())
        print_compiled_cost_analysis(compiled.cost_analysis())
        # Do not skip the first batch during training
        dataloader = itertools.chain([first_videos], dataloader)
    print(f"Starting training from step {step}...")
    while step < args.num_steps:
        for videos in dataloader:
            # --- Train step ---
            rng, _rng_mask = jax.random.split(rng, 2)
            inputs = dict(videos=videos, mask_rng=_rng_mask)
            loss, recon, metrics = train_step(optimizer, inputs)