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

            z_B1M = self.spatial_norm(z_B1M)
            z_B1M = self.spatial_attention(z_B1M)
            z_B1M = jnp.reshape(z_B1M, (B, M))
            z_BTNM = x_BTNM.at[:, pos_index[0], pos_index[1], :].set(z_B1M)
[tab] experiments/sample.sh (shellscript)
source .venv/bin/activate

data_dir="$PWD/data_arrayrecord/dummy"
ckpt_dir="$PWD/checkpoints/causal_dynamics_openai_grain_tok_restore"

export XLA_FLAGS='--xla_gpu_autotune_level=0'
export PYTHONUNBUFFERED=1
srun python sample.py \
    --dyna_type "causal" \
    --batch_size 1 \
    --seq_len 2 \
    --start_frame 1 \
    --checkpoint $ckpt_dir \
    --data_dir $data_dir
[content] /fast/home/franz.srambical/jafar/utils/nn.py (python)
        # Handle causal mask for cached decoder self-attention (from nnx.MultiHeadAttention)



        if mask is not None:


            mask_4d = _rearrange(mask)


            # NOTE: We need to broadcast T and S dimensions to target_seq_len since cudnn attention strictly checks the mask shape


            # https://github.com/jax-ml/jax/issues/28974


            # https://github.com/jax-ml/jax/blob/08c7677393672ccb85c10f1ed0bd506905c3c994/jax/_src/cudnn/fused_attention_stablehlo.py#L1830


            # https://github.com/jax-ml/jax/blob/08c7677393672ccb85c10f1ed0bd506905c3c994/jax/_src/cudnn/fused_attention_stablehlo.py#L337


            mask_4d = einops.repeat(mask_4d, "... 1 1 -> ... t s", t=target_seq_len, s=target_seq_len)


            mask_4d = mask_4d.astype(jnp.bool)


        else:


            mask_4d = attention_mask[jnp.newaxis, jnp.newaxis, :, :]  # (1, 1, seq_len, seq_len)
[content] /fast/home/franz.srambical/jafar/utils/nn.py (python)
    flax.nnx.MultiHeadAttention provides tensors with shape (batch..., length, num_heads, head_dim),
    but jax.nn.dot_product_attention expects (batch, length, num_heads, head_dim). We reshape to
    ensure compatibility. cuDNN's flash attention additionally requires a sequence length that
    is a multiple of 4. We pad the sequence length to the nearest multiple of 4 and mask
    accordingly. Note that cuDNN requires the mask to be broadcast before calling the attention
    function due to strict shape checking.
    """

    # FIXME (f.srambical): keys and values could have a different dimensionality
    def attention_fn(query_BSHD, key_BSHD, value_BSHD, bias=None, mask_B111=None, **kwargs):
        implementation = "cudnn" if use_flash_attention else None

        def _merge_batch_dims(x):
            return einops.rearrange(x, "... l h k -> (...) l h k")

        def _pad(x):
            return jnp.pad(x, ((0, 0), (0, pad_size), (0, 0), (0, 0)))

        original_shape = query_BSHD.shape
        original_seq_len = query_BSHD.shape[-3]

        # Pad to nearest multiple of 4
        T = ((original_seq_len + 3) // 4) * 4
        pad_size = T - original_seq_len

        query_BTHD = _pad(_merge_batch_dims(query_BSHD))
        key_BTHD = _pad(_merge_batch_dims(key_BSHD))
        value_BTHD = _pad(_merge_batch_dims(value_BSHD))
        B = query_BTHD.shape[0]

        attention_mask = jnp.ones((T, T), dtype=jnp.bool_)
        attention_mask = attention_mask.at[original_seq_len:, :].set(False)
        attention_mask = attention_mask.at[:, original_seq_len:].set(False)

        # Handle causal mask for cached decoder self-attention (from nnx.MultiHeadAttention)
        if mask_B111 is not None:
            mask_B111 = _merge_batch_dims(mask_B111)
            # We need to broadcast T and S dimensions to target_seq_len since cudnn attention strictly checks the mask shape
            # https://github.com/jax-ml/jax/issues/28974
            # https://github.com/jax-ml/jax/blob/08c7677393672ccb85c10f1ed0bd506905c3c994/jax/_src/cudnn/fused_attention_stablehlo.py#L1830
            # https://github.com/jax-ml/jax/blob/08c7677393672ccb85c10f1ed0bd506905c3c994/jax/_src/cudnn/fused_attention_stablehlo.py#L337
            mask_B1QK = einops.repeat(mask_B111, "... 1 1 -> ... t s", t=T, s=T)
            mask_B1QK = mask_B111.astype(jnp.bool)
        else:
            mask_11QK = attention_mask[jnp.newaxis, jnp.newaxis, :, :]
            mask_B1QK = jnp.broadcast_to(mask_11QK, (B, 1, T, T))

        bias_4d = _pad(_merge_batch_dims(bias)) if bias is not None else None

        # NOTE: jax.nn.dot_product_attention does not support dropout
        output_4d = jax.nn.dot_product_attention(
            query=query_BTHD,
            key=key_BTHD,
            value=value_BTHD,
            bias=bias_4d,
            mask=mask_B1QK,
[content] genie.py (python)

        for block in dynamics_causal.transformer.blocks:
            block.spatial_attention.init_cache((B * T, 1, self.dyna_dim), dtype=self.dtype)
            block.temporal_attention.init_cache((B * (N + 1), 1, self.dyna_dim), dtype=self.dtype)
[content] sample.py (python)

        for block in dynamics_causal.transformer.blocks:
            block.spatial_attention.init_cache((B * T, 1, self.dyna_dim), dtype=self.dtype)
            block.temporal_attention.init_cache((B * (N + 1), 1, self.dyna_dim), dtype=self.dtype)
[tab] utils/preprocess.py (python)
import einops
import jax
import jax.numpy as jnp


def patchify(videos: jax.Array, size: int) -> jax.Array:
    B, T, H, W, C = videos.shape
    x = jnp.pad(videos, ((0, 0), (0, 0), (0, -H % size), (0, -W % size), (0, 0)))
    return einops.rearrange(
        x, "b t (hn hp) (wn wp) c -> b t (hn wn) (hp wp c)", hp=size, wp=size
    )


def unpatchify(patches: jax.Array, size: int, h_out: int, w_out: int) -> jax.Array:
    h_pad = -h_out % size
    hn = (h_out + h_pad) // size
    x = einops.rearrange(
        patches,
        "b t (hn wn) (hp wp c) -> b t (hn hp) (wn wp) c",
        hp=size,
        wp=size,
        hn=hn,
    )
    return x[:, :, :h_out, :w_out]
[content] /fast/home/franz.srambical/jafar/genie.py (python)
            updated_token_idxs_BSN = final_carry_causal[1]
            new_carry = (rng, updated_token_idxs_BSN)
            return new_carry, None

        # --- Run the autoregressive generation using jax.lax.scan ---
        initial_carry = (batch["rng"], token_idxs_BSN)
        timesteps_to_scan = jnp.arange(T, seq_len)
        final_carry, _ = jax.lax.scan(
            generation_step_fn, initial_carry, timesteps_to_scan
        )
        final_token_idxs_BSN = final_carry[1]
[content] /fast/home/franz.srambical/jafar/genie.py (python)
        def generation_step_fn(
            carry: tuple[jax.Array, jax.Array], step_t: jax.Array
        ) -> tuple[tuple[jax.Array, jax.Array], None]:
            rng, current_token_idxs_BSN = carry