autogluon.tabular 1.3.2b20250712__py3-none-any.whl → 1.3.2b20250714__py3-none-any.whl

autogluon/tabular/models/mitra/_internal/models/tab2d.py

@@ -0,0 +1,667 @@
+from typing import Optional, Union
+
+import einops
+import einx
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from safetensors.torch import save_file
+from huggingface_hub import hf_hub_download
+from safetensors.torch import load_file
+import os
+import json
+
+# Try to import flash attention, but make it optional
+try:
+    from flash_attn.bert_padding import pad_input, unpad_input
+    from flash_attn.flash_attn_interface import flash_attn_varlen_func
+    FLASH_ATTN_AVAILABLE = True
+except ImportError:
+    FLASH_ATTN_AVAILABLE = False
+
+from torch.utils.checkpoint import checkpoint
+
+from ..._internal.config.enums import Task
+from ..._internal.models.base import BaseModel
+from ..._internal.models.embedding import (
+    Tab2DEmbeddingX, 
+    Tab2DEmbeddingYClasses, 
+    Tab2DEmbeddingYRegression,
+    Tab2DQuantileEmbeddingX,
+)
+
+
+class Tab2D(BaseModel):
+
+    def __init__(
+        self,
+        dim: int,
+        dim_output: int,
+        n_layers: int,
+        n_heads: int,
+        task: Union[str, Task],
+        use_pretrained_weights: bool,
+        path_to_weights: str,
+        device: str = "cuda",  # Add device parameter
+    ) -> None:
+
+        super().__init__()
+
+        self.dim = dim
+        self.dim_output = dim_output
+        self.n_layers = n_layers
+        self.n_heads = n_heads
+        self.task = task
+        self.device_type = device
+
+        # Determine if we can use flash attention
+        self.use_flash_attn = FLASH_ATTN_AVAILABLE and device.startswith("cuda")
+
+        if type(self.task) == str:
+            self.task = Task[self.task]
+
+        self.x_quantile = Tab2DQuantileEmbeddingX(dim)
+        self.x_embedding = Tab2DEmbeddingX(dim)
+
+
+        match self.task:
+            case Task.CLASSIFICATION:
+                self.y_embedding = Tab2DEmbeddingYClasses(dim, dim_output)     # type: nn.Module
+            case Task.REGRESSION:
+                if self.dim_output == 1:
+                    self.y_embedding = Tab2DEmbeddingYRegression(dim)
+                else:
+                    self.y_embedding = Tab2DEmbeddingYClasses(dim, dim_output)
+            case _:
+                raise ValueError(f"Task {task} not supported")
+
+        self.layers = nn.ModuleList()
+
+        for _ in range(n_layers):
+            self.layers.append(Layer(dim, n_heads, self.use_flash_attn))
+
+        self.final_layer_norm = nn.LayerNorm(dim)
+
+        self.final_layer = nn.Linear(dim, dim_output, bias=True)
+
+        if use_pretrained_weights:
+            if device == "cpu":
+                # For CPU, use weights_only=False since CUDA checkpoints are incompatible with weights_only=True
+                self.load_state_dict(torch.load(path_to_weights, weights_only=False, map_location=torch.device('cpu')))
+            else:
+                # For GPU, use weights_only=True for security
+                self.load_state_dict(torch.load(path_to_weights, weights_only=True, map_location=device))
+        else:
+            self.init_weights()
+
+
+    def forward(
+            self,
+            x_support: torch.Tensor, # (b, n_s, f)
+            y_support: torch.Tensor, # (b, n_s)
+            x_query: torch.Tensor, # (b, n_q, f)
+            padding_features: torch.Tensor, # (b, f), "1" represents padding, "0" represents valid values
+            padding_obs_support: torch.Tensor, # (b, n_s)
+            padding_obs_query__: torch.Tensor, # (b, n_q)
+        ):
+
+        """
+        x_support is (batch_size, n_observations_support, n_features)
+        y_support is (batch_size, n_observations_support)
+
+        x_query is (batch_size, n_observations_query, n_features)
+
+        returns:
+
+        y_query is (batch_size, n_observations_query, n_classes)
+
+        syntax:
+        b = batch size
+        s = number of observations
+        f = number of features
+        d = dimension of embedding
+        c = number of classes
+        """
+
+        x_query__ = x_query
+
+        batch_size = x_support.shape[0]
+        n_obs_support = x_support.shape[1]
+        n_obs_query__ = x_query__.shape[1]
+
+        x_support, x_query__ = self.x_quantile(x_support, x_query__, padding_obs_support, padding_features)
+        x_support = self.x_embedding(x_support) # (b, n_s, f, d)
+        x_query__ = self.x_embedding(x_query__) # (b, n_q, f, d)
+        y_support, y_query__ = self.y_embedding(y_support, padding_obs_support, n_obs_query__) # (b, n_s, 1, d), (b, n_q, 1, d)
+
+        support, pack_support = einops.pack((y_support, x_support), 'b s * d') # (b, n_s, f+1, d)
+        query__, pack_query__ = einops.pack((y_query__, x_query__), 'b s * d') # (b, n_q, f+1, d)
+
+        padding_features_y = torch.zeros((batch_size, 1), device=padding_features.device, dtype=torch.bool) # (b, 1)
+        padding_features, _ = einops.pack((padding_features_y, padding_features), 'b *') # (b, f+1)
+
+        if self.use_flash_attn:
+            padder_support = Padder(support, padding_obs_support, padding_features)
+            padder_query__ = Padder(query__, padding_obs_query__, padding_features)
+
+            support = padder_support.base_to_obs(support) # (n_valid_s, d)
+            query__ = padder_query__.base_to_obs(query__)  # (n_valid_q, d)
+
+            for layer in self.layers:
+                support, query__ = checkpoint(layer, support, query__, padder_support, padder_query__, use_reentrant=False) # (n_valid_s, d), (n_valid_q, d)
+
+            query__ = self.final_layer_norm(query__)
+            query__ = self.final_layer(query__) # (n_valid_q, d)
+
+            query__ = padder_query__.obs_to_base(query__) # (b, n_q, f+1, c)
+        else:
+            # For CPU/non-flash attention, work with standard tensor format
+            for layer in self.layers:
+                support, query__ = checkpoint(layer, support, query__, None, None,
+                                            batch_size, padding_obs_support, padding_obs_query__, padding_features, use_reentrant=False)
+
+            query__ = self.final_layer_norm(query__)
+            query__ = self.final_layer(query__) # (b, n_q, f+1, c)
+
+        y_query__, x_query__ = einops.unpack(query__, pack_query__, 'b s * c') # (b, n_q, 1, c), (b, n_q, f, c)
+
+        match self.task:
+            # output has shape (batch_size, n_observations_query, n_features, n_classes)
+            # we want to remove the n_features dimension, and for regression, the n_classes dimension
+            case Task.REGRESSION:
+                if self.dim_output == 1:
+                    y_query__ = y_query__[:, :, 0, 0]
+                else:
+                    y_query__ = y_query__[:, :, 0, :]
+            case Task.CLASSIFICATION:
+                y_query__ = y_query__[:, :, 0, :]
+            case _:
+                raise ValueError(f"Task {self.task} not supported")
+
+        return y_query__
+
+
+    def init_weights(self) -> None:
+
+        nn.init.normal_(self.x_embedding.x_embedding.weight, mean=0.0, std=1.0)
+        nn.init.normal_(self.x_embedding.x_embedding.bias, mean=0.0, std=1.0)
+        nn.init.normal_(self.y_embedding.y_embedding.weight, mean=0.0, std=1.0)
+        nn.init.normal_(self.y_embedding.y_mask.weight, mean=0.0, std=1.0)
+
+        # default PyTorch initialization for everything else
+
+
+    def save_pretrained(self, save_directory: str):
+        os.makedirs(save_directory, exist_ok=True)
+
+        save_file(self.state_dict(), os.path.join(save_directory, "model.safetensors"))
+
+        config = {
+            "dim": self.dim,
+            "dim_output": self.dim_output,
+            "n_layers": self.n_layers,
+            "n_heads": self.n_heads,
+            "task": str(self.task).upper(),
+        }
+        with open(os.path.join(save_directory, "config.json"), "w") as f:
+            json.dump(config, f)
+
+            
+    @classmethod
+    def from_pretrained(cls, path_or_repo_id: str, device: str = "cuda") -> "Tab2D":
+
+        config_path = hf_hub_download(repo_id=path_or_repo_id, filename="config.json")
+        with open(config_path, "r") as f:
+            config = json.load(f)
+
+        model = cls(
+            dim=config["dim"],
+            dim_output=config["dim_output"],
+            n_layers=config["n_layers"],
+            n_heads=config["n_heads"],
+            task=config["task"],
+            use_pretrained_weights=False,
+            path_to_weights="",
+            device=device
+        )
+
+        weights_path = hf_hub_download(repo_id=path_or_repo_id, filename="model.safetensors")
+        state_dict = load_file(weights_path, device=device)
+        model.load_state_dict(state_dict)
+
+        return model
+    
+    
+class Padder(torch.nn.Module):
+
+    def __init__(self, x: torch.Tensor, padding_mask: torch.Tensor, feature_mask: torch.Tensor) -> None:
+
+        super().__init__()
+
+        self.padding_mask = padding_mask
+        self.feature_mask = feature_mask
+
+        n_obs, n_feat = x.shape[1], x.shape[2]
+        self.batch_size = x.shape[0]
+
+        if not FLASH_ATTN_AVAILABLE:
+            # CPU fallback: implement simplified padding logic without flash attention
+            self._init_cpu_fallback(x, n_obs, n_feat)
+            return
+
+        # GPU path with flash attention
+        self._init_flash_attn(x, n_obs, n_feat)
+
+    def _init_cpu_fallback(self, x: torch.Tensor, n_obs: int, n_feat: int):
+        """Initialize CPU-compatible version without flash attention dependencies."""
+        # For CPU, we don't need the complex unpadding/padding logic
+        # We'll implement pass-through methods that preserve tensor shapes
+        self.cpu_mode = True
+
+        # Store original shapes for reference
+        self.original_shape = x.shape
+        self.n_obs = n_obs
+        self.n_feat = n_feat
+
+        # These attributes won't be used in CPU mode but need to exist for compatibility
+        self.cu_seqlens_o = None
+        self.cu_seqlens_f = None
+        self.cu_seqlens_fo = None
+        self.cu_seqlens_of = None
+        self.max_seqlen_in_batch_o = None
+        self.max_seqlen_in_batch_f = None
+        self.max_seqlen_in_batch_fo = None
+        self.max_seqlen_in_batch_of = None
+
+    def _init_flash_attn(self, x: torch.Tensor, n_obs: int, n_feat: int):
+        """Initialize GPU version with flash attention."""
+        self.cpu_mode = False
+
+        # Original flash attention initialization logic
+        x_o, self.indices_o, self.cu_seqlens_o, self.max_seqlen_in_batch_o = unpad_input(x, ~self.padding_mask)
+
+        self.feature_mask_big = einops.repeat(self.feature_mask, 'b f -> b s f', s=n_obs)
+        self.feature_mask_big, _, _, _ = unpad_input(self.feature_mask_big, ~self.padding_mask)
+        x_of, self.indices_of, self.cu_seqlens_of, self.max_seqlen_in_batch_of = unpad_input(x_o, ~self.feature_mask_big)
+
+        x_rearranged = einx.rearrange('b s f d -> b f s d', x)
+        x_f, self.indices_f, self.cu_seqlens_f, self.max_seqlen_in_batch_f = unpad_input(x_rearranged, ~self.feature_mask)
+
+        self.padding_mask_big = einops.repeat(self.padding_mask, 'b s -> b f s', f=n_feat)
+        self.padding_mask_big, _, _, _ = unpad_input(self.padding_mask_big, ~self.feature_mask)
+        x_fo, self.indices_fo, self.cu_seqlens_fo, self.max_seqlen_in_batch_fo = unpad_input(x_f, ~self.padding_mask_big)
+
+        self.batch_size_f = x_f.shape[0]
+        self.batch_size_o = x_o.shape[0]
+
+        t = torch.arange(self.indices_fo.shape[0]).unsqueeze(1).to(x.device)
+        self.obs_to_feat_indices = self.base_to_feat(self.obs_to_base(t)).squeeze(1)
+        self.feat_to_obs_indices = self.base_to_obs(self.feat_to_base(t)).squeeze(1)
+
+    def base_to_obs(self, x: torch.Tensor) -> torch.Tensor:
+        if hasattr(self, 'cpu_mode') and self.cpu_mode:
+            # CPU fallback: reshape for standard attention
+            # Convert from (b, s, f, d) to (b*s, f*d) or similar flattened format
+            b, s, f, d = x.shape
+            return x.view(b * s, f * d)
+
+        # GPU path with flash attention
+        x = einx.rearrange('b s f d -> b f s d', x)
+        x, _, _, _ = unpad_input(x, ~self.feature_mask)
+        x, _, _, _ = unpad_input(x, ~self.padding_mask_big)
+        return x
+
+    def base_to_feat(self, x: torch.Tensor) -> torch.Tensor:
+        if hasattr(self, 'cpu_mode') and self.cpu_mode:
+            # CPU fallback: reshape for standard attention
+            # Convert from (b, s, f, d) to (b*f, s*d) or similar flattened format
+            b, s, f, d = x.shape
+            return x.view(b * f, s * d)
+
+        # GPU path with flash attention
+        x, _, _, _ = unpad_input(x, ~self.padding_mask)
+        x, _, _, _ = unpad_input(x, ~self.feature_mask_big)
+        return x
+
+    def obs_to_base(self, x: torch.Tensor) -> torch.Tensor:
+        if hasattr(self, 'cpu_mode') and self.cpu_mode:
+            # CPU fallback: reshape back to base format
+            # This is the inverse of base_to_obs
+            total_elements = x.numel()
+            expected_d = self.original_shape[-1]  # last dimension
+            b, s, f = self.batch_size, self.n_obs, self.n_feat
+            return x.view(b, s, f, expected_d)
+
+        # GPU path with flash attention
+        x = pad_input(x, self.indices_fo, self.batch_size_f, self.max_seqlen_in_batch_fo)
+        x = pad_input(x, self.indices_f, self.batch_size, self.max_seqlen_in_batch_f)
+        x = einx.rearrange('b f s d -> b s f d', x)
+        return x
+
+    def feat_to_base(self, x: torch.Tensor) -> torch.Tensor:
+        if hasattr(self, 'cpu_mode') and self.cpu_mode:
+            # CPU fallback: reshape back to base format
+            # This is the inverse of base_to_feat
+            total_elements = x.numel()
+            expected_d = self.original_shape[-1]  # last dimension
+            b, s, f = self.batch_size, self.n_obs, self.n_feat
+            return x.view(b, s, f, expected_d)
+
+        # GPU path with flash attention
+        x = pad_input(x, self.indices_of, self.batch_size_o, self.max_seqlen_in_batch_of)
+        x = pad_input(x, self.indices_o, self.batch_size, self.max_seqlen_in_batch_o)
+        return x
+
+    def obs_to_feat(self, x: torch.Tensor) -> torch.Tensor:
+        if hasattr(self, 'cpu_mode') and self.cpu_mode:
+            # CPU fallback: simple pass-through or basic reshaping
+            return x
+
+        # GPU path with flash attention
+        x = x[self.obs_to_feat_indices]
+        return x
+
+    def feat_to_obs(self, x: torch.Tensor) -> torch.Tensor:
+        if hasattr(self, 'cpu_mode') and self.cpu_mode:
+            # CPU fallback: simple pass-through or basic reshaping
+            return x
+
+        # GPU path with flash attention
+        x = x[self.feat_to_obs_indices]
+        return x
+
+
+class Layer(torch.nn.Module):
+
+    def __init__(self, dim: int, n_heads: int, use_flash_attn: bool) -> None:
+
+        super().__init__()
+
+        self.layer_norm1 = nn.LayerNorm(dim)
+        self.attention1 = MultiheadAttention(dim, n_heads, use_flash_attn)
+        self.layer_norm2 = nn.LayerNorm(dim)
+        self.linear1 = nn.Linear(dim, dim*4, bias=True)
+        self.linear2 = nn.Linear(dim*4, dim, bias=True)
+
+        self.layer_norm3 = nn.LayerNorm(dim)
+        self.attention2 = MultiheadAttention(dim, n_heads, use_flash_attn)
+        self.layer_norm4 = nn.LayerNorm(dim)
+        self.linear3 = nn.Linear(dim, dim*4, bias=True)
+        self.linear4 = nn.Linear(dim*4, dim, bias=True)
+
+        self.use_flash_attn = use_flash_attn
+
+
+    def forward(
+            self,
+            support: torch.Tensor,
+            query__: torch.Tensor,
+            padder_support: Optional[Padder],
+            padder_query__: Optional[Padder],
+            batch_size: Optional[int] = None,
+            padding_obs_support: Optional[torch.Tensor] = None,
+            padding_obs_query__: Optional[torch.Tensor] = None,
+            padding_features: Optional[torch.Tensor] = None,
+        ) -> tuple[torch.Tensor, torch.Tensor]:
+
+        """
+        Input:
+        support in 'obs' format
+        query__ in 'obs' format
+
+        Output:
+        support in 'obs' format
+        query__ in 'obs' format
+        """
+
+        if self.use_flash_attn and padder_support is not None and padder_query__ is not None:
+            support_residual = support
+            query___residual = query__
+
+            support = self.layer_norm1(support)
+            query__ = self.layer_norm1(query__)
+
+            # attention across rows
+            support_att = self.attention1(
+                support, support, support,
+                cu_seqlens_q = padder_support.cu_seqlens_fo, max_seqlen_q = padder_support.max_seqlen_in_batch_fo,
+                cu_seqlens_k = padder_support.cu_seqlens_fo, max_seqlen_k = padder_support.max_seqlen_in_batch_fo
+            )
+            query___att = self.attention1(
+                query__, support, support,
+                cu_seqlens_q = padder_query__.cu_seqlens_fo, max_seqlen_q = padder_query__.max_seqlen_in_batch_fo,
+                cu_seqlens_k = padder_support.cu_seqlens_fo, max_seqlen_k = padder_support.max_seqlen_in_batch_fo
+            )
+
+            support = support_residual + support_att
+            query__ = query___residual + query___att
+
+            support_residual = support
+            query___residual = query__
+
+            support = self.layer_norm2(support)
+            query__ = self.layer_norm2(query__)
+
+            support = self.linear1(support)
+            query__ = self.linear1(query__)
+
+            support = torch.nn.functional.gelu(support)
+            query__ = torch.nn.functional.gelu(query__)
+
+            support = self.linear2(support)
+            query__ = self.linear2(query__)
+
+            support = support_residual + support
+            query__ = query___residual + query__
+
+            support = padder_support.obs_to_feat(support)
+            query__ = padder_query__.obs_to_feat(query__)
+
+            support_residual = support
+            query___residual = query__
+
+            support = self.layer_norm3(support)
+            query__ = self.layer_norm3(query__)
+
+            # attention across features
+            support = self.attention2(
+                support, support, support,
+                cu_seqlens_q = padder_support.cu_seqlens_of, max_seqlen_q = padder_support.max_seqlen_in_batch_of,
+                cu_seqlens_k = padder_support.cu_seqlens_of, max_seqlen_k = padder_support.max_seqlen_in_batch_of
+            )
+            query__ = self.attention2(
+                query__, query__, query__,
+                cu_seqlens_q = padder_query__.cu_seqlens_of, max_seqlen_q = padder_query__.max_seqlen_in_batch_of,
+                cu_seqlens_k = padder_query__.cu_seqlens_of, max_seqlen_k = padder_query__.max_seqlen_in_batch_of
+            )
+
+            support = support_residual + support
+            query__ = query___residual + query__
+
+            support_residual = support
+            query___residual = query__
+
+            support = self.layer_norm4(support)
+            query__ = self.layer_norm4(query__)
+
+            support = self.linear3(support)
+            query__ = self.linear3(query__)
+
+            support = torch.nn.functional.gelu(support)
+            query__ = torch.nn.functional.gelu(query__)
+
+            support = self.linear4(support)
+            query__ = self.linear4(query__)
+
+            support = support_residual + support
+            query__ = query___residual + query__
+
+            support = padder_support.feat_to_obs(support)
+            query__ = padder_query__.feat_to_obs(query__)
+
+            return support, query__
+        else:
+            # CPU/Standard attention path - ensure it matches the GPU logic exactly
+            # Input format: (b, s, f+1, d) where f+1 includes the y embedding
+            batch_size_actual, n_obs_support, n_feat_plus_one, dim = support.shape
+            _, n_obs_query, _, _ = query__.shape
+
+            if batch_size is None:
+                batch_size = batch_size_actual
+
+            # First attention block - attention across observations (rows)
+            support_residual = support
+            query___residual = query__
+
+            support = self.layer_norm1(support)
+            query__ = self.layer_norm1(query__)
+
+            # Reshape for row attention: (b, s, f+1, d) -> (b*(f+1), s, d)
+            support_flat = einops.rearrange(support, 'b s f d -> (b f) s d')
+            query___flat = einops.rearrange(query__, 'b s f d -> (b f) s d')
+
+            # attention across observations
+            support_att_flat = self.attention1(support_flat, support_flat, support_flat)
+            query___att_flat = self.attention1(query___flat, support_flat, support_flat)
+
+            # Reshape back: (b*(f+1), s, d) -> (b, s, f+1, d)
+            support_att = einops.rearrange(support_att_flat, '(b f) s d -> b s f d', b=batch_size)
+            query___att = einops.rearrange(query___att_flat, '(b f) s d -> b s f d', b=batch_size)
+
+            support = support_residual + support_att
+            query__ = query___residual + query___att
+
+            # First MLP block
+            support_residual = support
+            query___residual = query__
+
+            support = self.layer_norm2(support)
+            query__ = self.layer_norm2(query__)
+
+            support = self.linear1(support)
+            query__ = self.linear1(query__)
+
+            support = torch.nn.functional.gelu(support)
+            query__ = torch.nn.functional.gelu(query__)
+
+            support = self.linear2(support)
+            query__ = self.linear2(query__)
+
+            support = support_residual + support
+            query__ = query___residual + query__
+
+            # Second attention block - attention across features
+            support_residual = support
+            query___residual = query__
+
+            support = self.layer_norm3(support)
+            query__ = self.layer_norm3(query__)
+
+            # Reshape for feature attention: (b, s, f+1, d) -> (b*s, f+1, d)
+            support_feat = einops.rearrange(support, 'b s f d -> (b s) f d')
+            query___feat = einops.rearrange(query__, 'b s f d -> (b s) f d')
+
+            # attention across features
+            support_feat_att = self.attention2(support_feat, support_feat, support_feat)
+            query___feat_att = self.attention2(query___feat, query___feat, query___feat)
+
+            # Reshape back: (b*s, f+1, d) -> (b, s, f+1, d)
+            support_feat_att = einops.rearrange(support_feat_att, '(b s) f d -> b s f d', b=batch_size)
+            query___feat_att = einops.rearrange(query___feat_att, '(b s) f d -> b s f d', b=batch_size)
+
+            support = support_residual + support_feat_att
+            query__ = query___residual + query___feat_att
+
+            # Second MLP block
+            support_residual = support
+            query___residual = query__
+
+            support = self.layer_norm4(support)
+            query__ = self.layer_norm4(query__)
+
+            support = self.linear3(support)
+            query__ = self.linear3(query__)
+
+            support = torch.nn.functional.gelu(support)
+            query__ = torch.nn.functional.gelu(query__)
+
+            support = self.linear4(support)
+            query__ = self.linear4(query__)
+
+            support = support_residual + support
+            query__ = query___residual + query__
+
+            return support, query__
+
+
+class MultiheadAttention(torch.nn.Module):
+
+    def __init__(self, dim: int, n_heads: int, use_flash_attn: bool) -> None:
+
+        super().__init__()
+
+        self.use_flash_attn = use_flash_attn
+        self.dim = dim
+        self.n_heads = n_heads
+
+        self.q = nn.Linear(dim, dim, bias=True)
+        self.k = nn.Linear(dim, dim, bias=True)
+        self.v = nn.Linear(dim, dim, bias=True)
+        self.o = nn.Linear(dim, dim, bias=True)
+
+
+    def forward(
+            self,
+            query: torch.Tensor,
+            key: torch.Tensor,
+            value: torch.Tensor,
+            cu_seqlens_q: Optional[torch.Tensor] = None,
+            cu_seqlens_k: Optional[torch.Tensor] = None,
+            max_seqlen_q: Optional[int] = None,
+            max_seqlen_k: Optional[int] = None,
+        ) -> torch.Tensor:
+        """
+        b = batch size
+        s = number of observations
+        f = number of features
+        t = flashattention-compressed sequences of (batch, observations, features)
+        h = heads
+        d = dimension of embedding
+
+        input: (bsf, d) for flash attention or (b, s, d) for standard attention
+        output: (bsf, d) for flash attention or (b, s, d) for standard attention
+        """
+
+        q = self.q(query)
+        k = self.k(key)
+        v = self.v(value)
+
+        if self.use_flash_attn and cu_seqlens_q is not None:
+            q = einops.rearrange(q, 't (h d) -> t h d', h=self.n_heads) # (tokens, heads, dim), tokens is b*n*f w/o pad
+            k = einops.rearrange(k, 't (h d) -> t h d', h=self.n_heads)
+            v = einops.rearrange(v, 't (h d) -> t h d', h=self.n_heads)
+
+            output = flash_attn_varlen_func(
+                q = q,
+                k = k,
+                v = v,
+                cu_seqlens_q = cu_seqlens_q, # num_seq+1, either b*n (w/o pad)+1, or b*f (w/o pad)+1
+                cu_seqlens_k = cu_seqlens_k,
+                max_seqlen_q = max_seqlen_q, # max sequence length, either n or f
+                max_seqlen_k = max_seqlen_k,
+                deterministic=True,
+            )
+
+            output = einops.rearrange(output, 't h d -> t (h d)')
+        else:
+            # Standard scaled dot-product attention for CPU
+            q = einops.rearrange(q, 'b t (h d) -> b h t d', h=self.n_heads)
+            k = einops.rearrange(k, 'b t (h d) -> b h t d', h=self.n_heads)
+            v = einops.rearrange(v, 'b t (h d) -> b h t d', h=self.n_heads)
+
+            output = F.scaled_dot_product_attention(q, k, v)
+            output = einops.rearrange(output, 'b h t d -> b t (h d)')
+
+        output = self.o(output)
+
+        return output

autogluon/tabular/models/mitra/_internal/utils/set_seed.py

@@ -0,0 +1,15 @@
+import random
+import numpy as np
+import torch
+
+
+def set_seed(seed: int) -> None:
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed(seed)
+    torch.cuda.manual_seed_all(seed)
+
+def seed_worker(worker_id: int) -> None:
+    worker_seed = torch.initial_seed() % 2**32
+    set_seed(worker_seed)