p-doom/crowd-code-dataset-0.2 huggingface dataset:
[content] jasmine/models/tokenizer.py (python)

    Dimension keys:
        B: batch size
        T: sequence length
        N: number of patches per frame
        L: latent dimension
        D: B * T * N
        H: height
        W: width
        C: number of channels
        P: patch token dimension (patch_size^2 * C)
[content] jasmine/utils/nn.py (python)

        self.input_norm1 = nnx.LayerNorm(
            num_features=self.input_dim,
            param_dtype=self.param_dtype,
            dtype=self.param_dtype,  # layer norm in full precision
            rngs=rngs,
[content] jasmine/utils/nn.py (python)

        )
        self.input_norm2 = nnx.LayerNorm(
            num_features=self.model_dim,
            param_dtype=self.param_dtype,
            dtype=self.param_dtype,  # layer norm in full precision
            rngs=rngs,
        )
[content] jasmine/utils/nn.py (python)

        gate_spatial_BTM = einops.rearrange(gate_spatial_BTM, "b t m -> (b t) 1 m")
        shift_temporal_BTM = einops.repeat(shift_temporal_BTM, "b t m -> (tile b) t m", tile=N)
        scale_temporal_BTM = einops.repeat(scale_temporal_BTM, "b t m -> (tile b) t m", tile=N)
        gate_temporal_BTM = einops.repeat(gate_temporal_BTM, "b t m -> (tile b) t m", tile=N)
        shift_mlp_BTM = einops.rearrange(shift_mlp_BTM, "b t m -> b t 1 m")
        scale_mlp_BTM = einops.rearrange(scale_mlp_BTM, "b t m -> b t 1 m")
        gate_mlp_BTM = einops.rearrange(gate_mlp_BTM, "b t m -> b t 1 m")
[content] jasmine/genie.py (python)
    dyna_logits_BTNV, dyna_mask = self.dynamics(outputs)
        outputs["token_logits"] = dyna_logits_BTNV
        outputs["mask"] = dyna_mask
        mle_indices_BTN = jnp.argmax(outputs["token_logits"], axis=-1)
        outputs["recon"] = self.tokenizer.decode(mle_indices_BTN, (H, W))
        if action_indices_E is not None:
            outputs["lam_indices"] = action_indices_E
[content] jasmine/genie.py (python)

        """
        Autoregressively samples up to `seq_len` future frames, following Figure 8 of the Genie 2 paper (Kumar et al. (2025)).

        - Input frames are tokenized once.
        - Future frames are generated autoregressively in token space.
        - All frames are detokenized in a single pass.
        """
[content] jasmine/genie.py (python)

            B: batch size
            T: number of input (conditioning) frames
            N: number of patches per frame
            L: latent dimension
            S: sequence length
            H: height
            W: width