p-doom/crowd-code-dataset-0.2 huggingface dataset:
[content] jasmine/genie.py (python)
            currupted_latents_BSNL = (
                tok_latents_BSNL * t_corrupt_context
                + (1 - t_corrupt_context) * noise_context_BSNL
            )

            # Create a mask for the frames to update
            frame_mask = jnp.arange(seq_len) < frame_i
            frame_mask = frame_mask.reshape(1, -1, 1, 1)  # Broadcast to BSNL shape
            # Update only the frames that should be corrupted
            corrupted_tok_latents_BSNL = jnp.where(
                frame_mask, currupted_latents_BSNL, tok_latents_BSNL
[content] jasmine/genie.py (python)
<<<<<<< Updated upstream
            currupted_latents_BSNL = (
                latents_BSNL * t_corrupt_context
=======
            corrupted_latents_BSNL = (
                tok_latents_BSNL * t_corrupt_context
>>>>>>> Stashed changes
                + (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(
<<<<<<< Updated upstream
                frame_mask_1S11, currupted_latents_BSNL, latents_BSNL
=======
                frame_mask, corrupted_latents_BSNL, tok_latents_BSNL
>>>>>>> Stashed changes
[content] jasmine/train_dynamics_diffusion.py (python)
"""
Parts of the diffusion training, sampling, and DiT implementation are adapted from:
https://github.com/kvfrans/shortcut-models

We extend their approach with the diffusion-forcing objective and integrate several
design elements inspired by Dreamer 4 (https://arxiv.org/abs/2509.24527).
"""