[content] /hkfs/home/project/hk-project-p0023960/tum_cte0515/Projects/jasmine/train_dynamics.py (python)
        genie.eval()
        #(loss, (recon, metrics)) = dynamics_loss_fn(genie, inputs, training=False, pred_full_frame=False)
        (loss_full_frame, (recon_full_frame, metrics_full_frame)) = dynamics_loss_fn(genie, inputs, training=False, pred_full_frame=True)
        (loss, (recon, metrics)) = (loss_full_frame, (recon_full_frame, metrics_full_frame))
        return loss, recon, metrics, loss_full_frame, recon_full_frame, metrics_full_frame


    def calculate_validation_metrics(val_dataloader, genie, rng):
        step = 0
        loss_per_step = []
        metrics_per_step = []
        loss_full_frame_per_step = []
        metrics_full_frame_per_step = []
        for videos in val_dataloader:
            rng, _rng_mask = jax.random.split(rng, 2)
            inputs = dict(videos=videos, mask_rng=_rng_mask)
            loss, recon, metrics, loss_full_frame, recon_full_frame, metrics_full_frame = val_step(genie, inputs)
            loss_per_step.append(loss)
            metrics_per_step.append(metrics)
            loss_full_frame_per_step.append(loss_full_frame)
            metrics_full_frame_per_step.append(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_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_losses = {
            "val_loss": np.mean(loss_per_step),
            "val_loss_full_frame": np.mean(loss_full_frame_per_step)
        }
        val_metrics.update(val_metrics_full_frame)
        val_metrics.update(val_losses)
        return 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
    )
    if args.val_data_dir:
        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, mask_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, mask_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 ---
            if args.val_data_dir and step % args.val_interval == 0:
                rng, _rng_mask_val = jax.random.split(rng, 2)
                print(f"Calculating validation metrics...")
                params = nnx.state(optimizer.model)
                genie, rng = build_model(args, rng, params)