p-doom/crowd-code-dataset-0.2 huggingface dataset:
[content] utils/nn.py (python)
        implementation = 'cudnn' if use_flash_attention else None
        return jax.nn.dot_product_attention(
            query=query,
            key=key,
            value=value,
            bias=bias,
            mask=mask,
            implementation=implementation,
            is_causal=is_causal,
            **kwargs
        )
[tab] /hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/models/tokenizer.py (python)
from typing import Dict, Any, Tuple

import flax.linen as nn

from utils.preprocess import patchify, unpatchify
from utils.nn import STTransformer, VectorQuantizer


class TokenizerVQVAE(nn.Module):
    """ST-ViVit VQ-VAE"""

    in_dim: int
    model_dim: int
    latent_dim: int
    num_latents: int
    patch_size: int
    num_blocks: int
    num_heads: int
    dropout: float
    codebook_dropout: float
    use_flash_attention: bool

    def setup(self):
        self.encoder = STTransformer(
            self.model_dim,
            self.latent_dim,
            self.num_blocks,
            self.num_heads,
            self.dropout,
            use_flash_attention=self.use_flash_attention,
        )
        self.vq = VectorQuantizer(
            self.latent_dim,
            self.num_latents,
            self.codebook_dropout,
        )
        self.out_dim = self.in_dim * self.patch_size**2
        self.decoder = STTransformer(
            self.model_dim,
            self.out_dim,
            self.num_blocks,
            self.num_heads,
            self.dropout,
            use_flash_attention=self.use_flash_attention,
        )

    def __call__(self, batch: Dict[str, Any], training: bool = True) -> Dict[str, Any]:
        H, W = batch["videos"].shape[2:4]
        outputs = self.vq_encode(batch["videos"], training)
        recon = self.decoder(outputs["z_q"])  # (B, T, H_down * W_down, C)
        recon = nn.sigmoid(recon)
        outputs["recon"] = unpatchify(recon, self.patch_size, H, W)
        return outputs

    def vq_encode(self, videos: Any, training: bool = True) -> Dict[str, Any]:
        # --- Preprocess + encode ---
        B, T = videos.shape[:2]
        x = patchify(videos, self.patch_size)
        N = x.shape[2]
        x = self.encoder(x)  # (B, T, N, E)

        # --- Vector quantize ---
        x = x.reshape(B * T * N, self.latent_dim)
        z_q, z, emb, indices = self.vq(x, training)
        z_q = z_q.reshape(B, T, N, self.latent_dim)
        indices = indices.reshape(B, T, N)
        return dict(z_q=z_q, z=z, emb=emb, indices=indices)

    def decode(self, indices: Any, video_hw: Tuple[int, int]):
        z = self.vq.codebook[indices]
        recon = self.decoder(z)
        recon = nn.sigmoid(recon)
        return unpatchify(recon, self.patch_size, *video_hw)
[content] utils/nn.py (python)
        if use_flash_attention:
            # Check if tensors have 5D shape (B, S, K, num_heads, head_dim)
            # This indicates heads are already separated
            if (query.ndim == 5 and key.ndim == 5 and value.ndim == 5 and
                query.shape[-2:] == key.shape[-2:] == value.shape[-2:]):
                # Reshape to combine heads with feature dimension for cuDNN
                # (B, S, K, num_heads, head_dim) -> (B, S, K, num_heads * head_dim)
                query_reshaped = query.reshape(query.shape[:-2] + (-1,))
                key_reshaped = key.reshape(key.shape[:-2] + (-1,))
                value_reshaped = value.reshape(value.shape[:-2] + (-1,))
                
                try:
                    result = jax.nn.dot_product_attention(
                        query=query_reshaped,
                        key=key_reshaped,
                        value=value_reshaped,
                        bias=bias,
                        mask=mask,
                        implementation='cudnn',
                        is_causal=is_causal,
                        **kwargs
                    )
                    # Reshape back to original format
                    return result.reshape(query.shape)
                except Exception as e:
                    # Fall back to default implementation if cuDNN fails
                    print(f"Warning: cuDNN attention failed, falling back to default: {e}")
                    return jax.nn.dot_product_attention(
                        query=query,
                        key=key,
                        value=value,
                        bias=bias,
                        mask=mask,
                        implementation=None,
                        is_causal=is_causal,
                        **kwargs
                    )
            else:
                # Use default implementation for other shapes
                return jax.nn.dot_product_attention(
                    query=query,
                    key=key,
                    value=value,
                    bias=bias,
                    mask=mask,
                    implementation=None,
                    is_causal=is_causal,
                    **kwargs
                )
        else:
            return jax.nn.dot_product_attention(
                query=query,
                key=key,
                value=value,
                bias=bias,
                mask=mask,
                implementation=None,
                is_causal=is_causal,
                **kwargs
            )
    return attention_fn
[content] utils/nn.py (python)
        if use_flash_attention:
            # Check if tensors have 5D shape (B, S, K, num_heads, head_dim)
            # This indicates heads are already separated
            if (query.ndim == 5 and key.ndim == 5 and value.ndim == 5 and
                query.shape[-2:] == key.shape[-2:] == value.shape[-2:]):
                # Reshape to combine heads with feature dimension for cuDNN
                # (B, S, K, num_heads, head_dim) -> (B, S, K, num_heads * head_dim)
                query_reshaped = query.reshape(query.shape[:-2] + (-1,))
                key_reshaped = key.reshape(key.shape[:-2] + (-1,))
                value_reshaped = value.reshape(value.shape[:-2] + (-1,))
                
                result = jax.nn.dot_product_attention(
                    query=query_reshaped,
                    key=key_reshaped,
                    value=value_reshaped,
                    bias=bias,
                    mask=mask,
                    implementation='cudnn',
                    is_causal=is_causal,
                    **kwargs
                )
                # Reshape back to original format
                return result.reshape(query.shape)
            else:
                # Use default implementation for other shapes
                return jax.nn.dot_product_attention(
                    query=query,
                    key=key,
                    value=value,
                    bias=bias,
                    mask=mask,
                    implementation=None,
                    is_causal=is_causal,
                    **kwargs
                )
        else:
            return jax.nn.dot_product_attention(
                query=query,
                key=key,
                value=value,
                bias=bias,
                mask=mask,
                implementation=None,
                is_causal=is_causal,
                **kwargs
            )
    return attention_fn
[content] utils/nn.py (python)
        if use_flash_attention:
            # Check if tensors have 5D shape (B, S, K, num_heads, head_dim)
            # This indicates heads are already separated
            if (query.ndim == 5 and key.ndim == 5 and value.ndim == 5 and
                query.shape[-2:] == key.shape[-2:] == value.shape[-2:]):
                # Reshape to combine heads with feature dimension for cuDNN
                # (B, S, K, num_heads, head_dim) -> (B, S, K, num_heads * head_dim)
                query_reshaped = jnp.einsum('bsknh->bsk(nh)', query)
                key_reshaped = jnp.einsum('bsknh->bsk(nh)', key)
                value_reshaped = jnp.einsum('bsknh->bsk(nh)', value)
                
                result = jax.nn.dot_product_attention(
                    query=query_reshaped,
                    key=key_reshaped,
                    value=value_reshaped,
                    bias=bias,
                    mask=mask,
                    implementation='cudnn',
                    is_causal=is_causal,
                    **kwargs
                )
                # Reshape back to original format
                return jnp.einsum('bsk(nh)->bsknh', result, n=query.shape[-2])
            else:
                # Use default implementation for other shapes
                return jax.nn.dot_product_attention(
                    query=query,
                    key=key,
                    value=value,
                    bias=bias,
                    mask=mask,
                    implementation=None,
                    is_causal=is_causal,
                    **kwargs
                )
        else:
            return jax.nn.dot_product_attention(
                query=query,
                key=key,
                value=value,
                bias=bias,
                mask=mask,
                implementation=None,
                is_causal=is_causal,
                **kwargs
            )
    return attention_fn
[content] utils/nn.py (python)
for cuDNN
                # (B, S, K, num_heads, head_dim) -> (B, S, K, num_heads * head_dim)
                query_reshaped = query.reshape(query.shape[:-2] + (-1,))
                key_reshaped = key.reshape(key.shape[:-2] + (-1,))
                value_reshaped = value.re
[content] utils/nn.py (python)
(
                    query=query_reshaped,
                    key=key_reshaped,
                    value=value_reshaped,
                    bias=bias,
                    mask=mask,
                    implementati
[content] utils/nn.py (python)
None,
                    is_causal=is_causal,
                    **kwargs
                )
        else:
            return jax.nn.dot_product_attention(
                query=query,
                key=key,
                value=value,
                bias=bias,
                mask=mask,
[content] utils/nn.py (python)
        if use_flash_attention:
            # Check if tensors have 5D shape (B, S, K, num_heads, head_dim)
            # This indicates heads are already separated
            if (query.ndim == 5 and key.ndim == 5 and value.ndim == 5 and
                query.shape[-2:] == key.shape[-2:] == value.shape[-2:]):
                # Reshape to combine heads with feature dimension for cuDNN
                # (B, S, K, num_heads, head_dim) -> (B, S, K, num_heads * head_dim)
                query_reshaped = einops.rearrange(query, 'b s k n h -> b s k (n h)')
                key_reshaped = einops.rearrange(key, 'b s k n h -> b s k (n h)')
                value_reshaped = einops.rearrange(value, 'b s k n h -> b s k (n h)')
                
                result = jax.nn.dot_product_attention(
                    query=query_reshaped,
                    key=key_reshaped,
                    value=value_reshaped,
                    bias=bias,
                    mask=mask,
                    implementation='cudnn',
                    is_causal=is_causal,
                    **kwargs
                )
                # Reshape back to original format
                return einops.rearrange(result, 'b s k (n h) -> b s k n h', n=query.shape[-2])
            else:
                # Use default implementation for other shapes
                return jax.nn.dot_product_attention(
                    query=query,
                    key=key,
                    value=value,
                    bias=bias,
                    mask=mask,
                    implementation=None,
                    is_causal=is_causal,
                    **kwargs
                )
        else:
            return jax.nn.dot_product_attention(
                query=query,
                key=key,
                value=value,
                bias=bias,
                mask=mask,
                implementation=None,
                is_causal=is_causal,
                **kwargs
            )
    return attention_fn
[content] utils/nn.py (python)
        
        # Flax MultiHeadAttention provides tensors with shape (batch..., length, num_heads, head_dim)
        # jax.nn.dot_product_attention expects (batch, length, num_heads, head_dim) after _ensure_4d
        # We need to reshape to ensure compatibility
        
        # Get original shapes for later restoration
        query_shape = query.shape
        key_shape = key.shape
        value_shape = value.shape
        
        # Reshape to 4D: (batch, length, num_heads, head_dim)
        # For query: (batch..., length, num_heads, head_dim) -> (batch, length, num_heads, head_dim)
        batch_dims = query_shape[:-3]  # All dimensions except the last 3
        query_4d = query.reshape(-1, *query_shape[-3:])  # Flatten batch dims
        
        # For key and value: same reshaping
        key_4d = key.reshape(-1, *key_shape[-3:])
        value_4d = value.reshape(-1, *value_shape[-3:])
        
        # Handle bias and mask if provided
        bias_4d = None
        if bias is not None:
            bias_shape = bias.shape
            bias_4d = bias.reshape(-1, *bias_shape[-3:])
        
        mask_4d = None
        if mask is not None:
            mask_shape = mask.shape
            mask_4d = mask.reshape(-1, *mask_shape[-3:])
        
        # Call jax.nn.dot_product_attention with 4D tensors
        output_4d = jax.nn.dot_product_attention(
            query=query_4d,
            key=key_4d,
            value=value_4d,
            bias=bias_4d,
            mask=mask_4d,
[content] utils/nn.py (python)
        query_4d = einops.rearrange(query, '(b ...) l h d -> b l h d', b=batch_size)
        key_4d = einops.rearrange(key, '(b ...) l h d -> b l h d', b=batch_size)
        value_4d = einops.rearrange(value, '(b ...) l h d -> b l h d', b=batch_size)
[content] utils/nn.py (python)
        
        # Flax MultiHeadAttention provides tensors with shape (batch..., length, num_heads, head_dim)
        # jax.nn.dot_product_attention expects (batch, length, num_heads, head_dim) after _ensure_4d
        # We need to reshape to ensure compatibility
        
        # Get original shapes for later restoration
        query_shape = query.shape
        key_shape = key.shape
        value_shape = value.shape
        
        # Reshape to 4D: (batch, length, num_heads, head_dim)
        # For query: (batch..., length, num_heads, head_dim) -> (batch, length, num_heads, head_dim)
        batch_dims = query_shape[:-3]  # All dimensions except the last 3
        batch_size = int(jnp.prod(jnp.array(batch_dims)))
        
        query_4d = einops.rearrange(query, '(b ...) l h d -> b l h d', b=batch_size)
        key_4d = einops.rearrange(key, '(b ...) l h d -> b l h d', b=batch_size)
        value_4d = einops.rearrange(value, '(b ...) l h d -> b l h d', b=batch_size)
        
        # Handle bias and mask if provided
        bias_4d = None
        if bias is not None:
            bias_shape = bias.shape
            bias_batch_size = int(jnp.prod(jnp.array(bias_shape[:-3])))
            bias_4d = einops.rearrange(bias, '(b ...) l h d -> b l h d', b=bias_batch_size)
        
        mask_4d = None
        if mask is not None:
            mask_shape = mask.shape
            mask_batch_size = int(jnp.prod(jnp.array(mask_shape[:-3])))
            mask_4d = einops.rearrange(mask, '(b ...) l h d -> b l h d', b=mask_batch_size)
        
        # Call jax.nn.dot_product_attention with 4D tensors
        output_4d = jax.nn.dot_product_attention(
            query=query_4d,
            key=key_4d,
            value=value_4d,
            bias=bias_4d,
            mask=mask_4d,
[content] utils/nn.py (python)
        # Reshape to 4D: (batch, length, num_heads, head_dim)
        # Use einops to flatten batch dimensions without computing batch_size explicitly
        query_4d = einops.rearrange(query, '... l h d -> (...) l h d')
        key_4d = einops.rearrange(key, '... l h d -> (...) l h d')
        value_4d = einops.rearrange(value, '... l h d -> (...) l h d')
[content] utils/nn.py (python)
        batch_dims = query.shape[:-3]  # All dimensions except the last 3
        batch_size = int(jnp.prod(jnp.array(batch_dims)))
        
        query_4d = einops.rearrange(query, '(b ...) l h d -> b l h d', b=batch_size)
        key_4d = einops.rearrange(key, '(b ...) l h d -> b l h d', b=batch_size)
        value_4d = einops.rearrange(value, '(b ...) l h d -> b l h d', b=batch_size)
[content] utils/nn.py (python)
        
        # Flax MultiHeadAttention provides tensors with shape (batch..., length, num_heads, head_dim)
        # jax.nn.dot_product_attention expects (batch, length, num_heads, head_dim) after _ensure_4d
        # We need to reshape to ensure compatibility
        
        # Reshape to 4D: (batch, length, num_heads, head_dim)
        # Use einops to flatten batch dimensions without computing batch_size explicitly
        query_4d = einops.rearrange(query, '... l h d -> (...) l h d')
        key_4d = einops.rearrange(key, '... l h d -> (...) l h d')
        value_4d = einops.rearrange(value, '... l h d -> (...) l h d')
        
        # Handle bias and mask if provided
        bias_4d = None
        if bias is not None:
            bias_4d = einops.rearrange(bias, '... l h d -> (...) l h d')
        
        mask_4d = None
        if mask is not None:
            mask_4d = einops.rearrange(mask, '... l h d -> (...) l h d')
        
        output_4d = jax.nn.dot_product_attention(
            query=query_4d,
            key=key_4d,
            value=value_4d,
            bias=bias_4d,
            mask=mask_4d,
[content] utils/nn.py (python)
        
        # Flax MultiHeadAttention provides tensors with shape (batch..., length, num_heads, head_dim)
        # jax.nn.dot_product_attention expects (batch, length, num_heads, head_dim) after _ensure_4d
        # We need to reshape to ensure compatibility
        
        # Store original shape for restoration
        original_shape = query.shape
        
        # Reshape to 4D: (batch, length, num_heads, head_dim)
        # Use einops to flatten batch dimensions without computing batch_size explicitly
        query_4d = einops.rearrange(query, '... l h d -> (...) l h d')
        key_4d = einops.rearrange(key, '... l h d -> (...) l h d')
        value_4d = einops.rearrange(value, '... l h d -> (...) l h d')
        
        # Handle bias and mask if provided
        bias_4d = None
        if bias is not None:
            bias_4d = einops.rearrange(bias, '... l h d -> (...) l h d')
        
        mask_4d = None
        if mask is not None:
            mask_4d = einops.rearrange(mask, '... l h d -> (...) l h d')
        
        output_4d = jax.nn.dot_product_attention(
            query=query_4d,
            key=key_4d,
            value=value_4d,
            bias=bias_4d,
            mask=mask_4d,
[content] utils/nn.py (python)
            # Handle different mask shapes
            if mask.ndim == 2:
                # 2D mask: (seq_len, seq_len) - broadcast to 4D
                mask_4d = mask[None, None, :, :]  # (1, 1, seq_len, seq_len)
            elif mask.ndim == 4:
                # 4D mask: (batch..., seq_len, seq_len) - flatten batch dims
                mask_4d = einops.rearrange(mask, '... l1 l2 -> (...) l1 l2')
            else:
                # For other shapes, try to reshape to 4D
                mask_4d = mask.reshape(-1, *mask.shape[-2:])
[content] utils/nn.py (python)
        
        # Flax MultiHeadAttention provides tensors with shape (batch..., length, num_heads, head_dim)
        # jax.nn.dot_product_attention expects (batch, length, num_heads, head_dim) after _ensure_4d
        # We need to reshape to ensure compatibility
        
        # Store original shape for restoration
        original_shape = query.shape
        
        # Reshape to 4D: (batch, length, num_heads, head_dim)
        # Use einops to flatten batch dimensions without computing batch_size explicitly
        query_4d = einops.rearrange(query, '... l h d -> (...) l h d')
        key_4d = einops.rearrange(key, '... l h d -> (...) l h d')
        value_4d = einops.rearrange(value, '... l h d -> (...) l h d')
        
        # Handle bias and mask if provided
        bias_4d = None
        if bias is not None:
            bias_4d = einops.rearrange(bias, '... l h d -> (...) l h d')
        
        mask_4d = None
        if mask is not None:
            # Handle different mask shapes
            if mask.ndim == 2:
                # 2D mask: (seq_len, seq_len) - broadcast to 4D
                mask_4d = mask[None, None, :, :]  # (1, 1, seq_len, seq_len)
            elif mask.ndim == 4:
                # 4D mask: (batch..., seq_len, seq_len) - flatten batch dims
                mask_4d = einops.rearrange(mask, '... l1 l2 -> (...) l1 l2')
            else:
                # For other shapes, try to reshape to 4D
                mask_4d = mask.reshape(-1, *mask.shape[-2:])
        
        output_4d = jax.nn.dot_product_attention(
            query=query_4d,
            key=key_4d,
            value=value_4d,
            bias=bias_4d,
            mask=mask_4d,
[content] utils/nn.py (python)
        
        # Flax MultiHeadAttention provides tensors with shape (batch..., length, num_heads, head_dim)
        # jax.nn.dot_product_attention expects (batch, length, num_heads, head_dim) after _ensure_4d
        # We need to reshape to ensure compatibility
        
        # Store original shape for restoration
        original_shape = query.shape
        
        # Reshape to 4D: (batch, length, num_heads, head_dim)
        # Use einops to flatten batch dimensions without computing batch_size explicitly
        query_4d = einops.rearrange(query, '... l h d -> (...) l h d')
        key_4d = einops.rearrange(key, '... l h d -> (...) l h d')
        value_4d = einops.rearrange(value, '... l h d -> (...) l h d')
        
        # Handle bias and mask if provided
        bias_4d = None
        if bias is not None:
            bias_4d = einops.rearrange(bias, '... l h d -> (...) l h d')
        
        mask_4d = None
        if mask is not None:
            # Handle different mask shapes
            if mask.ndim == 2:
                # 2D mask: (seq_len, seq_len) - broadcast to 4D
                mask_4d = mask[None, None, :, :]  # (1, 1, seq_len, seq_len)
            elif mask.ndim == 4:
                # 4D mask: (batch..., seq_len, seq_len) - flatten batch dims
                mask_4d = einops.rearrange(mask, '... l1 l2 -> (...) l1 l2')
            else:
                # For other shapes, try to reshape to 4D
                mask_4d = mask.reshape(-1, *mask.shape[-2:])
            
            # Convert to boolean dtype as expected by jax.nn.dot_product_attention
            mask_4d = mask_4d.astype(jnp.bool_)
        
        output_4d = jax.nn.dot_product_attention(
            query=query_4d,
            key=key_4d,
            value=value_4d,
            bias=bias_4d,
            mask=mask_4d,
[content] /hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/utils/nn.py (python)
        
        # Flax MultiHeadAttention provides tensors with shape (batch..., length, num_heads, head_dim)
        # jax.nn.dot_product_attention expects (batch, length, num_heads, head_dim) after _ensure_4d
        # We need to reshape to ensure compatibility
        
        # Store original shape for restoration
        original_shape = query.shape
        
        query_4d = einops.rearrange(query, '... l h d -> (...) l h d')
        key_4d = einops.rearrange(key, '... l h d -> (...) l h d')
        value_4d = einops.rearrange(value, '... l h d -> (...) l h d')
        
        bias_4d = None
        if bias is not None:
            bias_4d = einops.rearrange(bias, '... l h d -> (...) l h d')
        
        mask_4d = None
        if mask is not None:
            # 2D mask: (seq_len, seq_len) - broadcast to 4D
            mask_4d = mask[jnp.newaxis, jnp.newaxis, :, :]  # (1, 1, seq_len, seq_len)
            
            # Convert to boolean dtype as expected by jax.nn.dot_product_attention
            mask_4d = mask_4d.astype(jnp.bool_)
        
        output_4d = jax.nn.dot_product_attention(
            query=query_4d,
            key=key_4d,
            value=value_4d,
            bias=bias_4d,
            mask=mask_4d,
            implementation=implementation,
            is_causal=is_causal,
            **kwargs
        )
        
        output = output_4d.reshape(original_shape)
        
        return output
[content] utils/nn.py (python)

        # Flax MultiHeadAttention provides tensors with shape (batch..., length, num_heads, head_dim)
        # jax.nn.dot_product_attention expects (batch, length, num_heads, head_dim) after _ensure_4d
        # We need to reshape to ensure compatibility