p-doom/crowd-code-dataset-0.2 huggingface dataset:
[tab] /fast/home/franz.srambical/jafar/models/dynamics_causal.py (python)
from typing import Dict

import jax
import jax.numpy as jnp
import flax.nnx as nnx

from utils.nn import STTransformer


class DynamicsCausal(nnx.Module):
    """Causal dynamics model"""

    def __init__(
        self,
        model_dim: int,
        ffn_dim: int,
        num_latents: int,
        latent_action_dim: int,
        num_blocks: int,
        num_heads: int,
        dropout: float,
        param_dtype: jnp.dtype,
        dtype: jnp.dtype,
        use_flash_attention: bool,
        decode: bool,
        rngs: nnx.Rngs,
    ):
        self.model_dim = model_dim
        self.ffn_dim = ffn_dim
        self.num_latents = num_latents
        self.latent_action_dim = latent_action_dim
        self.num_blocks = num_blocks
        self.num_heads = num_heads
        self.dropout = dropout
        self.param_dtype = param_dtype
        self.dtype = dtype
        self.use_flash_attention = use_flash_attention
        self.decode = decode

        self.transformer = STTransformer(
            self.model_dim,
            self.model_dim,
            self.ffn_dim,
            self.num_latents,
            self.num_blocks,
            self.num_heads,
            self.dropout,
            self.param_dtype,
            self.dtype,
            use_flash_attention=self.use_flash_attention,
            spatial_causal=True,
            decode=self.decode,
            rngs=rngs,
        )
        self.patch_embed = nnx.Embed(self.num_latents, self.model_dim, rngs=rngs)
        self.action_up = nnx.Linear(
            self.latent_action_dim,
            self.model_dim,
            param_dtype=self.param_dtype,
            dtype=self.dtype,
            rngs=rngs,
        )

    def __call__(
        self, batch: Dict[str, jax.Array], training: bool = True
    ) -> tuple[jax.Array, jax.Array | None]:
        vid_embed = self.patch_embed(batch["video_tokens"])
        act_embed = self.action_up(batch["latent_actions"])
        vid_embed += jnp.pad(act_embed, ((0, 0), (1, 0), (0, 0), (0, 0)))
        vid_embed_padded = jnp.pad(vid_embed, ((0, 0), (0, 0), (1, 0), (0, 0)))

        logits = self.transformer(vid_embed_padded)[:, :, :-1]

        mask = jnp.ones(vid_embed.shape[:-1])
        return logits, mask
[tab] models/dynamics_causal.py (python)
from typing import Dict

import jax
import jax.numpy as jnp
import flax.nnx as nnx

from utils.nn import STTransformer


class DynamicsCausal(nnx.Module):
    """Causal dynamics model"""

    def __init__(
        self,
        model_dim: int,
        ffn_dim: int,
        num_latents: int,
        latent_action_dim: int,
        num_blocks: int,
        num_heads: int,
        dropout: float,
        param_dtype: jnp.dtype,
        dtype: jnp.dtype,
        use_flash_attention: bool,
        decode: bool,
        rngs: nnx.Rngs,
    ):
        self.model_dim = model_dim
        self.ffn_dim = ffn_dim
        self.num_latents = num_latents
        self.latent_action_dim = latent_action_dim
        self.num_blocks = num_blocks
        self.num_heads = num_heads
        self.dropout = dropout
        self.param_dtype = param_dtype
        self.dtype = dtype
        self.use_flash_attention = use_flash_attention
        self.decode = decode

        self.transformer = STTransformer(
            self.model_dim,
            self.model_dim,
            self.ffn_dim,
            self.num_latents,
            self.num_blocks,
            self.num_heads,
            self.dropout,
            self.param_dtype,
            self.dtype,
            use_flash_attention=self.use_flash_attention,
            spatial_causal=True,
            decode=self.decode,
            rngs=rngs,
        )
        self.patch_embed = nnx.Embed(self.num_latents, self.model_dim, rngs=rngs)
        self.action_up = nnx.Linear(
            self.latent_action_dim,
            self.model_dim,
            param_dtype=self.param_dtype,
            dtype=self.dtype,
            rngs=rngs,
        )

    def __call__(
        self, batch: Dict[str, jax.Array], training: bool = True
    ) -> tuple[jax.Array, jax.Array | None]:
        vid_embed = self.patch_embed(batch["video_tokens"])
        act_embed = self.action_up(batch["latent_actions"])
        vid_embed += jnp.pad(act_embed, ((0, 0), (1, 0), (0, 0), (0, 0)))
        vid_embed_padded = jnp.pad(vid_embed, ((0, 0), (0, 0), (1, 0), (0, 0)))

        logits = self.transformer(vid_embed_padded)[:, :, :-1]

        mask = jnp.ones(vid_embed.shape[:-1])
        return logits, mask
[content] /fast/home/franz.srambical/jafar/jasmine.py (python)
            current_token_sequence = jax.lax.dynamic_slice(
                token_idxs_full, (0, t, 0), (B, 1, N)
            )

            dyna_inputs = {
                "video_tokens": current_token_sequence,
                "latent_actions": action_tokens,
            }
            # The model will output logits for all patches in the sequence (which is just one frame).
            # FIXME: do we need the model to output logits for a single frame in order to use kv caching?
[tab] models/dynamics_maskgit.py (python)
from typing import Dict

import jax
import jax.numpy as jnp
import flax.nnx as nnx

from utils.nn import STTransformer


class DynamicsMaskGIT(nnx.Module):
    """MaskGIT dynamics model"""

    def __init__(
        self,
        model_dim: int,
        ffn_dim: int,
        num_latents: int,
        latent_action_dim: int,
        num_blocks: int,
        num_heads: int,
        dropout: float,
        mask_limit: float,
        param_dtype: jnp.dtype,
        dtype: jnp.dtype,
        use_flash_attention: bool,
        rngs: nnx.Rngs,
    ):
        self.model_dim = model_dim
        self.ffn_dim = ffn_dim
        self.num_latents = num_latents
        self.latent_action_dim = latent_action_dim
        self.num_blocks = num_blocks
        self.num_heads = num_heads
        self.dropout = dropout
        self.mask_limit = mask_limit
        self.param_dtype = param_dtype
        self.dtype = dtype
        self.use_flash_attention = use_flash_attention

        self.transformer = STTransformer(
            self.model_dim,
            self.model_dim,
            self.ffn_dim,
            self.num_latents,
            self.num_blocks,
            self.num_heads,
            self.dropout,
            self.param_dtype,
            self.dtype,
            use_flash_attention=self.use_flash_attention,
            spatial_causal=False,
            decode=False,
            rngs=rngs,
        )
        self.patch_embed = nnx.Embed(self.num_latents, self.model_dim, rngs=rngs)
        self.mask_token = nnx.Param(
            nnx.initializers.lecun_uniform()(rngs.params(), (1, 1, 1, self.model_dim))
        )
        self.action_up = nnx.Linear(
            self.latent_action_dim,
            self.model_dim,
            param_dtype=self.param_dtype,
            dtype=self.dtype,
            rngs=rngs,
        )

    def __call__(
        self, batch: Dict[str, jax.Array], training: bool = True
    ) -> tuple[jax.Array, jax.Array | None]:
        # --- Mask videos ---
        vid_embed = self.patch_embed(batch["video_tokens"])
        if training:
            batch_size = vid_embed.shape[0]
            _rng_prob, *_rngs_mask = jax.random.split(batch["mask_rng"], batch_size + 1)
            mask_prob = jax.random.uniform(
                _rng_prob, shape=(batch_size,), minval=self.mask_limit
            )
            per_sample_shape = vid_embed.shape[1:-1]
            mask = jax.vmap(
                lambda rng, prob: jax.random.bernoulli(rng, prob, per_sample_shape),
                in_axes=(0, 0),
            )(jnp.asarray(_rngs_mask), mask_prob)
            mask = mask.at[:, 0].set(False)
            vid_embed = jnp.where(
                jnp.expand_dims(mask, -1), self.mask_token.value, vid_embed
            )
        else:
            mask = None

        # --- Predict transition ---
        act_embed = self.action_up(batch["latent_actions"])
        vid_embed += jnp.pad(act_embed, ((0, 0), (1, 0), (0, 0), (0, 0)))
        logits = self.transformer(vid_embed)
        return logits, mask
[content] /fast/home/franz.srambical/jafar/models/dynamics_causal.py (python)
        # breakpoint()
        act_embed = self.action_up(batch["latent_actions"])
        # FIXME: this is the culprit!
        # vid_embed += jnp.pad(act_embed, ((0, 0), (1, 0), (0, 0), (0, 0)))
        vid_embed += act_embed
        vid_embed_padded = jnp.pad(vid_embed, ((0, 0), (0, 0), (1, 0), (0, 0)))
        # breakpoint()
[content] /fast/home/franz.srambical/jafar/utils/nn.py (python)
            
            # If mask has more than 2 dimensions, reduce it to 2D
            if mask_bool.ndim > 2:
                # For masks with shape (batch, seq_len, ...), take the first element along extra dims
                mask_2d = mask_bool[..., 0, 0] if mask_bool.ndim >= 4 else mask_bool[..., 0]
            else:
                mask_2d = mask_bool
            
            # Pad the 2D mask
[content] /fast/home/franz.srambical/jafar/utils/nn.py (python)
            
            # Debug: print mask shape
            print(f"DEBUG: mask_bool.shape = {mask_bool.shape}, ndim = {mask_bool.ndim}")
            
            # If mask has more than 2 dimensions, reduce it to 2D
            if mask_bool.ndim > 2:
                # For masks with shape (batch, seq_len, ...), take the first element along extra dims
                mask_2d = mask_bool[..., 0, 0] if mask_bool.ndim >= 4 else mask_bool[..., 0]
            else:
                mask_2d = mask_bool
            
            print(f"DEBUG: mask_2d.shape = {mask_2d.shape}, ndim = {mask_2d.ndim}")
            
            # Pad the 2D mask
[content] /fast/home/franz.srambical/jafar/utils/nn.py (python)
            
            # Debug: print mask shape
            print(f"DEBUG: mask_bool.shape = {mask_bool.shape}, ndim = {mask_bool.ndim}")
            
            # Reduce mask to 2D by taking the first element along all extra dimensions
            if mask_bool.ndim > 2:
                # For any mask with more than 2 dims, squeeze out all single dimensions
                mask_2d = mask_bool.squeeze()
                # If it's still not 2D, take the first element along remaining extra dims
                if mask_2d.ndim > 2:
                    # Take first element along all extra dimensions
                    indices = [0] * (mask_2d.ndim - 2)
                    mask_2d = mask_2d[..., *indices]
            else:
                mask_2d = mask_bool
            
            print(f"DEBUG: mask_2d.shape = {mask_2d.shape}, ndim = {mask_2d.ndim}")
            
            # Pad the 2D mask
[content] /fast/home/franz.srambical/jafar/utils/nn.py (python)
            
            # Debug: print mask shape
            print(f"DEBUG: mask_bool.shape = {mask_bool.shape}, ndim = {mask_bool.ndim}")
            
            # Reduce mask to 2D by taking the first element along all extra dimensions
            if mask_bool.ndim > 2:
                # For any mask with more than 2 dims, squeeze out all single dimensions
                mask_2d = mask_bool.squeeze()
                # If it's still not 2D, take the first element along remaining extra dims
                if mask_2d.ndim > 2:
                    # Take first element along all extra dimensions
                    for _ in range(mask_2d.ndim - 2):
                        mask_2d = mask_2d[..., 0]
            else:
                mask_2d = mask_bool
            
            print(f"DEBUG: mask_2d.shape = {mask_2d.shape}, ndim = {mask_2d.ndim}")
            
            # Pad the 2D mask
[content] /fast/home/franz.srambical/jafar/utils/nn.py (python)
            
            # Debug: print mask shape
            print(f"DEBUG: mask_bool.shape = {mask_bool.shape}, ndim = {mask_bool.ndim}")
            
            # Reduce mask to 2D by taking the first element along all extra dimensions
            if mask_bool.ndim > 2:
                # For any mask with more than 2 dims, squeeze out all single dimensions
                mask_2d = mask_bool.squeeze()
                # If it's 1D, reshape to 2D (assuming it's a sequence mask)
                if mask_2d.ndim == 1:
                    mask_2d = mask_2d.reshape(1, -1)
                # If it's still not 2D, take the first element along remaining extra dims
                elif mask_2d.ndim > 2:
                    # Take first element along all extra dimensions
                    for _ in range(mask_2d.ndim - 2):
                        mask_2d = mask_2d[..., 0]
            else:
                mask_2d = mask_bool
            
            print(f"DEBUG: mask_2d.shape = {mask_2d.shape}, ndim = {mask_2d.ndim}")
            
            # Pad the 2D mask
[content] /fast/home/franz.srambical/jafar/utils/nn.py (python)
            
            # Reduce mask to 2D by taking the first element along all extra dimensions
            if mask_bool.ndim > 2:
                # For any mask with more than 2 dims, squeeze out all single dimensions
                mask_2d = mask_bool.squeeze()
                # If it's 1D, reshape to 2D (assuming it's a sequence mask)
                if mask_2d.ndim == 1:
                    mask_2d = mask_2d.reshape(1, -1)
                # If it's still not 2D, take the first element along remaining extra dims
                elif mask_2d.ndim > 2:
                    # Take first element along all extra dimensions
                    for _ in range(mask_2d.ndim - 2):
                        mask_2d = mask_2d[..., 0]
            else:
                mask_2d = mask_bool
            
            # Pad the 2D mask