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

            self.dynamics = DynamicsCausal(
                model_dim=self.dyna_dim,
                ffn_dim=self.dyna_ffn_dim,
                num_latents=self.num_patch_latents,
                latent_action_dim=self.latent_action_dim,
            )
[content] genie.py (python)
num_blocks=self.dyna_num_blocks,
                num_heads=self.dyna_num_heads,
                dropout=self.dropout,
                param_dtype=self.param_dtype,
                dtype=self.dtype,
                use_flash_attention=self.use_flash_attention,
                rngs=rngs,
[content] models/tokenizer.py (python)
Dimension key:

B: batch size
L: sequence length
M: memory length (length of sequence being attended to)
D: model dimension (sometimes called d_model or embedding_dim)
V: vocabulary size
F: feed-forward subnetwork hidden size
H: number of attention heads in a layer
K: size of each attention key or value (sometimes called d_kv)
"""
[content] models/tokenizer.py (python)

        """
        Dimension key:

        B: batch size
        L: sequence length
        M: memory length (length of sequence being attended to)
        D: model dimension (sometimes called d_model or embedding_dim)
        V: vocabulary size
        F: feed-forward subnetwork hidden size
        H: number of attention heads in a layer
        K: size of each attention key or value (sometimes called d_kv)
        """
[content] models/tokenizer.py (python)

        Dimension key:

        B: batch size
        T: sequence length
        N: number of patches per frame
        L: latent dimension
        P: (patch size)^2
        H: number of attention heads in a layer
        K: size of each attention key or value (sometimes called d_kv)
[content] models/tokenizer.py (python)
    def decode(self, token_indices: jax.Array, video_height_width: Tuple[int, int]) -> jax.Array:
        # Lookup embeddings from codebook using token indices
        codebook_embeddings = self.vq.codebook[token_indices]
        
        # Decode embeddings back to patch representations
        decoded_patches = self.decoder(codebook_embeddings)
        
        # Apply sigmoid activation and convert to float32 for numerical stability
        decoded_patches = decoded_patches.astype(jnp.float32)
        decoded_patches = nnx.sigmoid(decoded_patches)
        decoded_patches = decoded_patches.astype(self.dtype)
        
        # Reconstruct full video from patches
        reconstructed_video = unpatchify(decoded_patches, self.patch_size, *video_height_width)
        
        return reconstructed_video
[content] models/tokenizer.py (python)
    def decode(self, indices: jax.Array, video_hw: Tuple[int, int]) -> jax.Array:
        z = self.vq.codebook[indices]
        recon = self.decoder(z)
        recon = recon.astype(jnp.float32)
        recon = nnx.sigmoid(recon)
        recon = recon.astype(self.dtype)
        return unpatchify(recon, self.patch_size, *video_hw)
[content] models/tokenizer.py (python)
    def decode(self, indices_BTN: jax.Array, video_hw: Tuple[int, int]) -> jax.Array:
        z_BTNL = self.vq.codebook[indices_BTN]
        recon_BTNP = self.decoder(z_BTNL)
        recon_BTNP = recon_BTNP.astype(jnp.float32)
        recon_BTNP = nnx.sigmoid(recon_BTNP)
        recon_BTNP = recon_BTNP.astype(self.dtype)
        return unpatchify(recon_BTNP, self.patch_size, *video_hw)
[content] models/tokenizer.py (python)
        x_BTNL_flat = x_BTNL.reshape(B * T * N, self.latent_dim)
        z_q_flat, z_flat, emb, indices_flat = self.vq(x_BTNL_flat, training)
        z_q_BTNL = z_q_flat.reshape(B, T, N, self.latent_dim)
        indices_BTN = indices_flat.reshape(B, T, N)
        return dict(z_q=z_q_BTNL, z=z_flat, emb=emb, indices=indices_BTN)
[content] models/tokenizer.py (python)
        x = x_BTNL.reshape(B * T * N, self.latent_dim)
        z_q, z, emb, indices = self.vq(x, training)
        z_q_BTNL = z_q.reshape(B, T, N, self.latent_dim)
        indices_BTN = indices.reshape(B, T, N)
        return dict(z_q=z_q_BTNL, z=z, emb=emb, indices=indices_BTN)
[content] models/lam.py (python)

    B: batch size
    T: sequence length
    N: number of patches per frame
    L: latent dimension
    P: (patch size)^2
    H: number of attention heads in a layer
    D: B*T*N
    L: latent dimension
    H: height
    W: width
    C: number of channels
    O: output dimension (P*C)