p-doom/crowd-code-dataset-0.2 huggingface dataset:
[content] jasmine/models/genie.py (python)
        action_tokens_BSm11L = jnp.reshape(action_tokens_EL, (B, seq_len - 1, 1, -1))
        act_embed_BS1L = jnp.pad(action_tokens_BSm11L, ((0, 0), (1, 0), (0, 0), (0, 0)))
        act_embed_BS1L = self.dynamics.action_up(act_embed_BS1L)

        t_corrupt_context = 1 - diffusion_corrupt_context_factor
        t_corrupt_context = jnp.argmin(
            jnp.abs(jnp.arange(diffusion_steps) / diffusion_steps - t_corrupt_context)
        )
        t_corrupt_context = t_corrupt_context / diffusion_steps

        @nnx.scan(in_axes=(nnx.Carry, 0), out_axes=nnx.Carry)
        def denoise_step_fn(carry, denoise_step: jax.Array) -> jax.Array:
            latents_BSNL, frame_i, rng = carry
            rng, _rng_noise_context = jax.random.split(rng)

            # We need to reconstruct the submodule inside scan body to prevent trace context mismatches
            dynamics_diffusion = DynamicsDiffusion(
                model_dim=self.dyna_dim,
                ffn_dim=self.dyna_ffn_dim,
                latent_patch_dim=self.latent_patch_dim,
                latent_action_dim=self.latent_action_dim,
                num_blocks=self.dyna_num_blocks,
                num_heads=self.dyna_num_heads,
                denoise_steps=self.diffusion_denoise_steps,
                dropout=self.dropout,
                param_dtype=self.param_dtype,
                dtype=self.dtype,
                use_flash_attention=self.use_flash_attention,
                rngs=nnx.Rngs(0),
                decode=self.decode,
            )
            nnx.update(dynamics_diffusion, dynamics_state)

            # corrupt the context frames like in Dreamer 4 section 3.2
            denoise_step_BS = jnp.ones((B, seq_len)) * diffusion_steps - 1
            denoise_step_BS = denoise_step_BS.at[:, frame_i].set(denoise_step)

            noise_context_BSNL = jax.random.normal(
                _rng_noise_context, (B, seq_len, N, L)
            )
            corrupted_latents_BSNL = (
                latents_BSNL * t_corrupt_context
                + (1 - t_corrupt_context) * noise_context_BSNL
            )
            frame_mask = jnp.arange(seq_len) < frame_i
            frame_mask_1S11 = frame_mask.reshape(1, seq_len, 1, 1)
            corrupted_tok_latents_BSNL = jnp.where(
                frame_mask_1S11, corrupted_latents_BSNL, latents_BSNL
            )

            # --- Predict transition ---
            denoise_step_embed_BS1L = self.dynamics.timestep_embed(
                denoise_step_BS
            ).reshape(B, seq_len, 1, self.latent_patch_dim)
            inputs_BSNp2L = jnp.concatenate(
                [act_embed_BS1L, denoise_step_embed_BS1L, corrupted_tok_latents_BSNL],
                axis=2,
            )
            pred_latents_BSNL = dynamics_diffusion.diffusion_transformer(
                inputs_BSNp2L,
            )
            pred_latents_BSNL = pred_latents_BSNL[:, :, 2:]
[content] jasmine/models/genie.py (python)
        ctx_signal_level = 1 - diffusion_corrupt_context_factor
        ctx_signal_level = jnp.argmin(
            jnp.abs(jnp.arange(diffusion_steps) / diffusion_steps - ctx_signal_level)
        )
        ctx_signal_level = ctx_signal_level / diffusion_steps

        @nnx.scan(in_axes=(nnx.Carry, 0), out_axes=nnx.Carry)
        def denoise_step_fn(
            carry: tuple[jax.Array, jax.Array, jax.Array], denoise_step: jax.Array
        ) -> tuple[jax.Array, jax.Array, jax.Array]:
            latents_BSNL, frame_i, rng = carry
            rng, _rng_noise_context = jax.random.split(rng)

            # We need to reconstruct the submodule inside scan body to prevent trace context mismatches
            dynamics_diffusion = DynamicsDiffusion(
                model_dim=self.dyna_dim,
                ffn_dim=self.dyna_ffn_dim,
                latent_patch_dim=self.latent_patch_dim,
                latent_action_dim=self.latent_action_dim,
                num_blocks=self.dyna_num_blocks,
                num_heads=self.dyna_num_heads,
                denoise_steps=self.diffusion_denoise_steps,
                dropout=self.dropout,
                param_dtype=self.param_dtype,
                dtype=self.dtype,
                use_flash_attention=self.use_flash_attention,
                rngs=nnx.Rngs(0),
                decode=self.decode,
            )
            nnx.update(dynamics_diffusion, dynamics_state)

            # corrupt the context frames like in Dreamer 4 section 3.2
            denoise_step_BS = jnp.ones((B, seq_len)) * diffusion_steps - 1
            denoise_step_BS = denoise_step_BS.at[:, frame_i].set(denoise_step)

            noise_context_BSNL = jax.random.normal(
                _rng_noise_context, (B, seq_len, N, L)
            )
            corrupted_latents_BSNL = (
                latents_BSNL * ctx_signal_level
                + (1 - ctx_signal_level) * noise_context_BSNL
            )
            frame_mask = jnp.arange(seq_len) < frame_i
            frame_mask_1S11 = frame_mask.reshape(1, seq_len, 1, 1)
            corrupted_tok_latents_BSNL = jnp.where(
                frame_mask_1S11, corrupted_latents_BSNL, latents_BSNL
            )

            # --- Predict transition ---
            action_tokens_BSm11L = jnp.reshape(
                action_tokens_EL, (B, seq_len - 1, 1, -1)
            )
            act_embed_BSm11L = self.dynamics.action_up(action_tokens_BSm11L)
            act_embed_BS1L = jnp.pad(act_embed_BSm11L, ((0, 0), (1, 0), (0, 0), (0, 0)))

            denoise_step_embed_BS1L = self.dynamics.timestep_embed(
                denoise_step_BS
            ).reshape(B, seq_len, 1, self.latent_patch_dim)