[content] /hkfs/home/project/hk-project-p0023960/tum_cte0515/Projects/jasmine/models/dynamics.py (python)
   178|        if pred_full_frame:
   179|            # --- Extract submodule states ---
   180|            patch_embed_state = nnx.state(self.patch_embed)
   181|            action_up_state = nnx.state(self.action_up)
   182|            transformer_state = nnx.state(self.transformer)
   183|            model_dim = self.model_dim
   184|            latent_action_dim = self.latent_action_dim
   185|            num_latents = self.num_latents
   186|            ffn_dim = self.ffn_dim
   187|            num_blocks = self.num_blocks
   188|            num_heads = self.num_heads
   189|            dropout = self.dropout
   190|            param_dtype = self.param_dtype
   191|            dtype = self.dtype
   192|            use_flash_attention = self.use_flash_attention
   193|            decode = True  # For full frame prediction
   194|
   195|            def _pred_full_frame(carry, step_n):
   196|                video_tokens_BTN, final_logits_BTNV = carry
   197|                # Reconstruct submodules inside scan body
   198|                patch_embed = nnx.Embed(num_latents, model_dim)
   199|                nnx.update(patch_embed, patch_embed_state)
   200|                action_up = nnx.Linear(
   201|                    latent_action_dim, model_dim, param_dtype=param_dtype, dtype=dtype
   202|                )
   203|                nnx.update(action_up, action_up_state)
   204|                transformer = Transformer(
   205|                    model_dim, model_dim, ffn_dim, num_latents, num_blocks, num_heads,
   206|                    dropout, param_dtype, dtype, use_flash_attention=use_flash_attention,
   207|                    decode=decode
   208|                )
   209|                nnx.update(transformer, transformer_state)
   210|
   211|                vid_embed_BTNM = patch_embed(video_tokens_BTN)
   212|                act_embed_BTm11M = action_up(latent_actions_BTm11L)
   213|                padded_act_embed_BT1M = jnp.pad(
   214|                    act_embed_BTm11M, ((0, 0), (1, 0), (0, 0), (0, 0))
   215|                )
   216|                vid_embed_BTNp1M = jnp.concatenate(
   217|                    [padded_act_embed_BT1M, vid_embed_BTNM], axis=2
   218|                )
   219|                step_logits_BTNp1V = transformer(vid_embed_BTNp1M)
   220|                step_logits_BV = step_logits_BTNp1V[:, -1, step_n, :]
   221|                final_logits_BTNV = final_logits_BTNV.at[:, -1, step_n].set(step_logits_BV)
   222|                sampled_token_idxs_B = jnp.argmax(step_logits_BV, axis=-1)
   223|                video_tokens_BTN = video_tokens_BTN.at[:, -1, step_n].set(
   224|                    sampled_token_idxs_B
   225|                )
   226|                return (video_tokens_BTN, final_logits_BTNV), None
   227|
   228|            (_, final_logits_BTNV), _ = jax.lax.scan(
   229|                _pred_full_frame,
   230|                (video_tokens_BTN, jnp.zeros((
   231|                    video_tokens_BTN.shape[0],
   232|                    video_tokens_BTN.shape[1],
   233|                    video_tokens_BTN.shape[2],
   234|                    self.num_latents))),
   235|                jnp.arange(video_tokens_BTN.shape[1])
   236|            )
   237|            mask_out = jnp.zeros_like(video_tokens_BTN)
   238|            mask_out = mask_out.at[:, -1].set(True)
   239|            return final_logits_BTNV, mask_out