[content] /hkfs/home/project/hk-project-p0023960/tum_cte0515/Projects/jasmine/train_dynamics.py (python)
            rng, _rng_mask = jax.random.split(rng, 2)
            inputs = dict(videos=videos, rng=_rng_mask)
            val_outputs = val_step(genie, inputs)
            loss_per_step.append(val_outputs["loss"])
            metrics_per_step.append(val_outputs["metrics"])
            if args.eval_full_frame:
                loss_full_frame_per_step.append(val_outputs["loss_full_frame"])
                metrics_full_frame_per_step.append(val_outputs["metrics_full_frame"])
            step += 1
            if step > args.val_steps:
                break

        if step < args.val_steps:
            print(
                f"Warning: Your validation dataset is too small to make val_steps many steps. Made {step} steps, expected {args.val_steps}"
            )

        val_metrics = {
            f"val_{key}": np.mean([float(m[key]) for m in metrics_per_step])
            for key in metrics_per_step[0].keys()
        }
        val_metrics["val_loss"] = np.mean(loss_per_step)
        if args.eval_full_frame:
            val_metrics_full_frame = {
                f"val_full_frame_{key}": np.mean(
                    [float(m[key]) for m in metrics_full_frame_per_step]
                )
                for key in metrics_full_frame_per_step[0].keys()
            }
            val_metrics.update(val_metrics_full_frame)
            val_metrics["val_loss_full_frame"] = np.mean(loss_full_frame_per_step)
        return val_dataloader, val_metrics, inputs, recon, recon_full_frame

    # --- TRAIN LOOP ---
    dataloader_train = (
        jax.make_array_from_process_local_data(videos_sharding, elem)
        for elem in train_iterator
    )
    dataloader_val = None
    if val_iterator:
        dataloader_val = (
            jax.make_array_from_process_local_data(videos_sharding, elem)
            for elem in val_iterator
        )
    if jax.process_index() == 0:
        first_videos = next(dataloader_train)
        sample_inputs = dict(videos=first_videos, rng=rng)
        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_train = itertools.chain([first_videos], dataloader_train)
    print(f"Starting training from step {step}...")
    first_step = step
    while step < args.num_steps:
        for videos in dataloader_train:
            # --- Train step ---
            rng, _rng_mask = jax.random.split(rng, 2)
            inputs = dict(videos=videos, rng=_rng_mask)
            loss, recon, metrics = train_step(optimizer, inputs)
            if step == first_step:
                print_mem_stats("After params initialized")
            metrics["lr"] = lr_schedule(step)
            print(f"Step {step}, loss: {loss}")
            step += 1

            # --- Validation loss ---
            val_results = {}
            if dataloader_val and step % args.val_interval == 0:
                rng, _rng_mask_val = jax.random.split(rng, 2)
                print("Calculating validation metrics...")
                (
                    dataloader_val,
                    val_metrics,
                    val_gt_batch,
                    val_recon,
                    val_recon_full_frame,
                ) = calculate_validation_metrics(
                    dataloader_val, optimizer.model, _rng_mask_val
                )
                print(f"Step {step}, validation loss: {val_metrics['val_loss']}")
                val_results = {
                    "metrics": val_metrics,
                    "gt_batch": val_gt_batch,
                    "recon": val_recon,
                    "full_frame": val_recon_full_frame,
                }

            # --- Logging ---
            if args.log:
                if step % args.log_interval == 0 and jax.process_index() == 0:
                    log_dict = {"loss": loss, "step": step, **metrics}
                    if val_results:
                        log_dict.update(val_results["metrics"])
                    wandb.log(log_dict)
                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 val_results:
                        val_results["gt_seq_val"] = (
                            val_results["gt_batch"]["videos"][0].astype(jnp.float32)
                            / 255.0
                        )
                        val_results["recon_seq_val"] = val_results["recon"].clip(0, 1)
                        val_comparison_seq = jnp.concatenate(
                            (val_results["gt_seq_val"], val_results["recon_seq_val"]),
                            axis=1,
                        )
                        val_results["val_comparison_seq"] = einops.rearrange(
                            val_comparison_seq * 255, "t h w c -> h (t w) c"
                        )
                        if args.eval_full_frame:
                            val_results["full_frame_seq_val"] = val_results[
                                "full_frame"
                            ][0].clip(0, 1)
                            val_results["val_full_frame_comparison_seq"] = (
                                jnp.concatenate(
                                    (
                                        val_results["gt_seq_val"],
                                        val_results["full_frame_seq_val"],
                                    ),
                                    axis=1,
                                )
                            )
                            val_results["val_full_frame_comparison_seq"] = (
                                einops.rearrange(
                                    val_results["val_full_frame_comparison_seq"] * 255,
                                    "t h w c -> h (t w) c",
                                )
                            )
                    # NOTE: Process-dependent control flow deliberately happens
                    # after indexing operation since it must not contain code
                    # sections that lead to cross-accelerator communication.
                    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))
                            ),
                        )
                        if val_results:
                            log_images.update(
                                dict(
                                    val_image=wandb.Image(
                                        np.asarray(
                                            val_results["gt_seq_val"][args.seq_len - 1]
                                        )
                                    ),
                                    val_recon=wandb.Image(
                                        np.asarray(
                                            val_results["recon_seq_val"][
                                                args.seq_len - 1
                                            ]
                                        )
                                    ),
                                    val_true_vs_recon=wandb.Image(
                                        np.asarray(
                                            val_results["val_comparison_seq"].astype(
                                                np.uint8
                                            )
                                        )
                                    ),
                                )
                            )
                            if args.eval_full_frame:
                                log_images.update(
                                    dict(
                                        val_full_frame=wandb.Image(
                                            np.asarray(
                                                val_results["full_frame_seq_val"][
                                                    args.seq_len - 1
                                                ]
                                            )
                                        ),
                                        val_true_vs_full_frame=wandb.Image(
                                            np.asarray(
                                                val_results[
                                                    "val_full_frame_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)
                if val_iterator: