suq 0.1.0__py3-none-any.whl

suq/SUQ_MLP.py

@@ -0,0 +1,10 @@
+from .diag_suq_mlp import SUQ_MLP_Diag
+
+def streamline_mlp(model, posterior, covariance_structure, likelihood, scale_init = 1.0):
+    if covariance_structure == 'diag':
+        return SUQ_MLP_Diag(org_model = model, 
+                            posterior_variance = posterior, 
+                            likelihood = likelihood,
+                            scale_init = scale_init)
+    else:
+        raise NotImplementedError(f"Covariance structure '{covariance_structure}' is not implemented.")

suq/SUQ_ViT.py

@@ -0,0 +1,14 @@
+from .diag_suq_transformer import SUQ_ViT_Diag
+
+def streamline_vit(model, posterior, covariance_structure, likelihood, MLP_deterministic, Attn_deterministic, attention_diag_cov, num_det_blocks, scale_init = 1.0):
+    if covariance_structure == 'diag':
+        return SUQ_ViT_Diag(ViT = model, 
+                            posterior_variance = posterior, 
+                            MLP_determinstic = MLP_deterministic, 
+                            Attn_determinstic = Attn_deterministic, 
+                            likelihood = likelihood,
+                            attention_diag_cov = attention_diag_cov, 
+                            num_det_blocks = num_det_blocks,
+                            scale_init = scale_init)
+    else:
+        raise NotImplementedError(f"Covariance structure '{covariance_structure}' is not implemented.")

suq/__init__.py

@@ -0,0 +1,4 @@
+from .SUQ_MLP import streamline_mlp
+from .SUQ_ViT import streamline_vit
+
+__all__ = ["streamline_mlp", "streamline_vit"]

suq/base_suq.py

@@ -0,0 +1,181 @@
+import torch
+import torch.nn as nn
+import numpy as np
+from tqdm import tqdm
+from torch.distributions import Categorical
+from torch.distributions.normal import Normal
+from torch.utils.data import DataLoader
+
+from suq.utils.utils import torch_dataset
+
+device = 'cuda' if torch.cuda.is_available() else 'cpu'
+
+class SUQ_Base(nn.Module):
+    """
+    Base class for SUQ models.
+
+    Provides core functionality for:
+    - Managing likelihood type (regression or classification)
+    - Probit-based approximation for classification
+    - NLPD-based fitting of the scale factor
+
+    Inputs:
+        likelihood (str): Either 'classification' or 'regression'
+        scale_init (float): Initial value for the scale factor parameter
+    """
+    
+    def __init__(self, likelihood, scale_init):
+        super().__init__()
+        
+        if likelihood not in ['classification', 'regression']:
+            raise ValueError(f"Invalid likelihood type {likelihood}")
+        
+        self.likelihood = likelihood
+        self.scale_factor = nn.Parameter(torch.Tensor([scale_init]).to(device))
+    
+    def probit_approximation(self, out_mean, out_var):
+        """
+        Applies a probit approximation to compute class probabilities from the latent Gaussian distribution.
+        
+        Inputs:
+            out_mean (Tensor): Latent function mean, shape [B, C]
+            out_var (Tensor): Latent function variance, shape [B, C] or [B, C, C]
+
+        Outputs:
+            posterior_predict_mean (Tensor): Predicted class probabilities, shape [B, C]
+        """
+
+        if out_var.dim() == 3:
+            kappa = 1 / torch.sqrt(1. + np.pi / 8 * out_var.diagonal(dim1=1, dim2=2))
+        else:
+            kappa = 1 / torch.sqrt(1. + np.pi / 8 * out_var)
+            
+        posterior_predict_mean = torch.softmax(kappa * out_mean, dim=-1)
+        return posterior_predict_mean
+
+    def fit_scale_factor(self, data_loader, n_epoches, lr, speedup = True, verbose = False):
+        """
+        Fits the scale factor for predictive variance using negative log predictive density (NLPD).
+
+        Inputs:
+            data_loader (DataLoader): Dataloader containing (input, target) pairs
+            n_epoches (int): Number of epochs for optimization
+            lr (float): Learning rate for scale optimizer
+            speedup (bool): If True (classification only), caches forward pass outputs to accelerate fitting
+            verbose (bool): If True, prints NLPD at each epoch
+
+        Outputs:
+            total_train_nlpd (List[float]): Average NLPD per epoch over training data
+        """
+        print("fit scale factor")
+        optimizer = torch.optim.Adam([self.scale_factor], lr)
+        total_train_nlpd = []
+        
+        # store intermediate result and pack it into a data loader, so we only need to do one forward pass
+        if speedup:
+            
+            if self.likelihood == 'regression':
+                raise ValueError(f"Speed up not supported for regression atm")
+            
+            if self.likelihood == 'classification':
+
+                f_mean = []
+                f_var = []
+                labels = []
+
+                for (X, y) in tqdm(data_loader, desc= "packing f_mean f_var into a dataloader"):
+                    out_mean, out_var = self.forward_latent(X.to(device))
+                    f_mean.append(out_mean.detach().cpu().numpy())
+                    f_var.append(out_var.detach().cpu().numpy())
+                    if y.dim() == 2:
+                        labels.append(y.numpy().argmax(1).reshape(-1, 1))
+                    if y.dim() == 1:
+                        labels.append(y.numpy().reshape(-1, 1))
+
+                f_mean = np.vstack(f_mean)
+                f_var = np.vstack(f_var)
+                labels = np.vstack(labels)
+
+                scale_fit_dataset = torch_dataset(f_mean, f_var, labels)
+                scale_fit_dataloader = DataLoader(scale_fit_dataset, batch_size=16, shuffle=True)
+
+                for epoch in tqdm(range(n_epoches), desc="fitting scaling factor"):
+                    running_nlpd = 0
+                    for data_pair in scale_fit_dataloader:
+                        x_mean, x_var_label = data_pair
+                        num_class = x_mean.shape[1]
+                        x_mean = x_mean.to(device)
+                        x_var, label = x_var_label.split(num_class, dim=1)
+                        x_var = x_var.to(device)
+                        label = label.to(device)
+
+                        optimizer.zero_grad()
+                        # make prediction
+                        x_var = x_var / self.scale_factor.to(device)
+                        posterior_predict_mean = self.probit_approximation(x_mean, x_var)
+                        # construct log posterior predictive distribution
+                        posterior_predictive_dist = Categorical(posterior_predict_mean)
+                        # calculate nlpd and update
+                        nlpd = -posterior_predictive_dist.log_prob(label).mean()
+                        nlpd.backward()
+                        optimizer.step()
+                        # log nlpd
+                        running_nlpd += nlpd.item()
+                    total_train_nlpd.append(running_nlpd / len(scale_fit_dataloader))
+                    if verbose:
+                        print(f"epoch {epoch + 1}, nlpd {total_train_nlpd[-1]}")
+                        
+                del scale_fit_dataloader
+                del scale_fit_dataset
+        
+        else:
+            
+            if self.likelihood == 'classification':
+                for epoch in tqdm(range(n_epoches), desc="fitting scaling factor"):
+                    running_nlpd = 0
+                    for (data, label) in data_loader:
+                        
+                        data = data.to(device)
+                        label = label.to(device)
+
+                        optimizer.zero_grad()
+                        # make prediction
+                        posterior_predict_mean = self.forward(data)
+                        # construct log posterior predictive distribution
+                        posterior_predictive_dist = Categorical(posterior_predict_mean)
+                        # calculate nlpd and update
+                        nlpd = -posterior_predictive_dist.log_prob(label).mean()
+                        nlpd.backward()
+                        optimizer.step()
+                        # log nlpd
+                        running_nlpd += nlpd.item()
+                    total_train_nlpd.append(running_nlpd / len(data_loader))
+                    if verbose:
+                        print(f"epoch {epoch + 1}, nlpd {total_train_nlpd[-1]}")
+
+
+            if self.likelihood == 'regression':
+                for epoch in tqdm(range(n_epoches), desc="fitting scaling factor"):
+                    running_nlpd = 0
+                    for (data, label) in data_loader:
+                        data = data.to(device)
+                        label = label.to(device)
+                        
+                        optimizer.zero_grad()
+                        # make prediction
+                        posterior_predict_mean, posterior_predict_var = self.forward(data)
+                        # construct log posterior predictive distribution
+                        posterior_predictive_dist = Normal(posterior_predict_mean, posterior_predict_var.sqrt())
+                        # calculate nlpd and update
+                        nlpd = -posterior_predictive_dist.log_prob(label).mean()
+                        nlpd.backward()
+                        optimizer.step()
+                        # log nlpd
+                        running_nlpd += nlpd.item()
+                
+                    total_train_nlpd.append(running_nlpd / len(data_loader))
+                    
+                    if verbose:
+                        print(f"epoch {epoch + 1}, nlpd {total_train_nlpd[-1]}")
+                    
+        return total_train_nlpd

suq/diag_suq_mlp.py

@@ -0,0 +1,308 @@
+import torch.nn as nn
+import torch.nn.functional as F
+import torch
+from torch.nn.utils import parameters_to_vector
+import numpy as np
+import copy
+
+from suq.base_suq import SUQ_Base
+
+def forward_aW_diag(a_mean, a_var, weight, bias, w_var, b_var):
+    """
+    compute mean and covariance of h = a @ W^T + b when posterior has diag covariance
+    
+    ----- Input -----
+    a_mean: [N, D_in] mean(a)
+    a_var: [N, D_in] a_var[i] = var(a_i)
+    weight: [D_out, D_in] W
+    bias: [D_out, ] b
+    b_var: [D_out, ] b_var[k]: var(b_k)
+    w_var: [D_out, D_in] w_cov[k][i]: var(w_ki)
+    ----- Output -----
+    h_mean: [N, D_out]
+    h_var: [N, D_out] h_var[k] = var(h_k)
+    """
+    
+    # calculate mean(h)
+    h_mean = F.linear(a_mean, weight, bias)
+    
+    # calculate var(h)
+    weight_mean2_var_sum = weight ** 2 + w_var # [D_out, D_in]
+    h_var = a_mean **2 @ w_var.T + a_var @ weight_mean2_var_sum.T + b_var
+    
+    return h_mean, h_var
+
+
+def forward_activation_implicit_diag(activation_func, h_mean, h_var):
+    """
+    given h ~ N(h_mean, h_cov), g(·), where h_cov is a diagonal matrix,
+    approximate the distribution of a = g(h) as 
+    a ~ N(g(h_mean), g'(h_mean)^T h_var g'(h_mean))
+    
+    input
+    activation_func: g(·)
+    h_mean: [N, D]
+    h_var: [N, D], h_var[i] = var(h_i)
+    
+    output
+    a_mean: [N, D]
+    a_var: [N, D]
+    """
+
+    h_mean_grad = h_mean.detach().clone().requires_grad_()
+    
+    a_mean = activation_func(h_mean_grad)
+    a_mean.retain_grad()
+    a_mean.backward(torch.ones_like(a_mean)) #[N, D]
+    
+    nabla = h_mean_grad.grad #[N, D]
+    a_var = nabla ** 2 * h_var
+    
+    return a_mean.detach(), a_var
+
+def forward_batch_norm_diag(h_var, bn_weight, bn_running_var, bn_eps):
+    """
+    Pass a distribution with diagonal covariance through BatchNorm layer
+    
+    Input
+        h_mean: mean of input distribution [B, D]
+        h_var: variance of input distribution [B, D]
+        bn_weight: batch norm scale factor [D, ]
+        bn_running_var: batch norm running variance [D, ]
+        bn_eps: batch norm eps
+    
+    Output
+        output_var [B, T, D]
+    """
+
+    scale_factor = (1 / (bn_running_var.reshape(1, -1) + bn_eps)) * bn_weight.reshape(1, -1) **2 # [B, D]
+    output_var = scale_factor * h_var # [B, D]
+    
+    return output_var
+
+class SUQ_Linear_Diag(nn.Module):
+    """
+    Linear layer with uncertainty propagation under SUQ, with a diagonal Gaussian posterior.
+    
+    Wraps a standard `nn.Linear` layer and applies closed-form mean and variance propagation. See the SUQ paper for theoretical background and assumptions.
+
+    Inputs:
+        org_linear (nn.Linear): The original linear layer to wrap
+        w_var (Tensor): Weight variances, shape [D_out, D_in]
+        b_var (Tensor): Bias variances, shape [D_out]
+    """
+    def __init__(self, org_linear, w_var, b_var):
+        super().__init__()
+        
+        self.weight = org_linear.weight.data
+        self.bias = org_linear.bias.data
+        self.w_var = w_var
+        self.b_var = b_var
+    
+    def forward(self, a_mean, a_var): 
+        """
+        Inputs:
+            a_mean (Tensor): Input mean, shape [N, D_in]
+            a_var (Tensor): Input variance, shape [N, D_in]
+
+        Outputs:
+            h_mean (Tensor): Output mean, shape [N, D_out]
+            h_var (Tensor): Output variance, shape [N, D_out]
+        """
+        
+        if a_var == None:
+            a_var = torch.zeros_like(a_mean).to(a_mean.device)
+            
+        h_mean, h_var = forward_aW_diag(a_mean, a_var, self.weight, self.bias, self.w_var, self.b_var)
+        
+        return h_mean, h_var
+
+class SUQ_Activation_Diag(nn.Module):
+    """
+    Activation layer with closed-form uncertainty propagation under SUQ, with a diagonal Gaussian posterior.
+
+    Wraps a standard activation function and applies a first-order approximation to propagate input variance through the nonlinearity. See the SUQ paper for theoretical background and assumptions.
+
+    Inputs:
+        afun (Callable): A PyTorch activation function (e.g. nn.ReLU())
+    """
+    
+    def __init__(self, afun):        
+        super().__init__()
+        self.afun = afun
+    
+    def forward(self, h_mean, h_var):
+        """
+        Inputs:
+            h_mean (Tensor): Input mean before activation, shape [N, D]
+            h_var (Tensor): Input variance before activation, shape [N, D]
+
+        Outputs:
+            a_mean (Tensor): Activated output mean, shape [N, D]
+            a_var (Tensor): Approximated output variance, shape [N, D]
+        """
+        a_mean, a_var = forward_activation_implicit_diag(self.afun, h_mean, h_var)
+        return a_mean, a_var
+
+class SUQ_BatchNorm_Diag(nn.Module):
+    """
+    BatchNorm layer with closed-form uncertainty propagation under SUQ, with a diagonal Gaussian posterior.
+
+    Wraps `nn.BatchNorm1d` and adjusts input variance using batch normalization statistics and scale parameters. See the SUQ paper for theoretical background and assumptions.
+
+    Inputs:
+        BatchNorm (nn.BatchNorm1d): The original batch norm layer
+    """
+    
+    def __init__(self, BatchNorm):
+        super().__init__()
+        
+        self.BatchNorm = BatchNorm
+    
+    def forward(self, x_mean, x_var):
+        """
+        Inputs:
+            x_mean (Tensor): Input mean, shape [B, D]
+            x_var (Tensor): Input variance, shape [B, D]
+
+        Outputs:
+            out_mean (Tensor): Output mean after batch normalization, shape [B, D]
+            out_var (Tensor): Output variance after batch normalization, shape [B, D]
+        """
+        
+        with torch.no_grad():
+        
+            out_mean = self.BatchNorm.forward(x_mean)
+            out_var = forward_batch_norm_diag(x_mean, x_var, self.BatchNorm.weight, 1e-5)
+            
+        return out_mean, out_var
+
+
+class SUQ_MLP_Diag(SUQ_Base):
+    """
+    Multilayer perceptron model with closed-form uncertainty propagation under SUQ, with a diagonal Gaussian posterior.
+
+    Wraps a standard MLP, converting its layers into SUQ-compatible components.
+    Supports both classification and regression via predictive Gaussian approximation.
+    
+    Note:
+        The input model should correspond to the latent function only:
+        - For regression, this is the full model (including final output layer).
+        - For classification, exclude the softmax layer and pass only the logit-producing part.
+
+    Inputs:
+        org_model (nn.Module): The original MLP model to convert
+        posterior_variance (Tensor): Flattened posterior variance vector
+        likelihood (str): Either 'classification' or 'regression'
+        scale_init (float, optional): Initial scale factor
+        sigma_noise (float, optional): noise level (for regression)
+    """
+    
+    def __init__(self, org_model, posterior_variance, likelihood, scale_init = 1.0, sigma_noise = None):
+        super().__init__(likelihood, scale_init)
+
+        self.sigma_noise = sigma_noise
+        self.convert_model(org_model, posterior_variance)
+    
+    def forward_latent(self, data, out_var = None):
+        """
+        Compute the predictive mean and variance of the latent function before applying the likelihood.
+
+        Traverses the model layer by layer, propagating mean and variance through each SUQ-wrapped layer.
+
+        Inputs:
+            data (Tensor): Input data, shape [B, D]
+            out_var (Tensor or None): Optional input variance, shape [B, D]
+
+        Outputs:
+            out_mean (Tensor): Output mean after final layer, shape [B, D_out]
+            out_var (Tensor): Output variance after final layer, shape [B, D_out]
+        """
+        
+        out_mean = data
+        
+        if isinstance(self.model, nn.Sequential):
+            for layer in self.model:
+                out_mean, out_var = layer.forward(out_mean, out_var)
+        ##TODO: other type of model            
+
+        out_var = out_var / self.scale_factor
+        
+        return out_mean, out_var
+    
+    def forward(self, data):
+        """
+        Compute the predictive distribution based on the model's likelihood setting.
+
+        For classification, use probit-approximation.
+        For regression, returns the latent mean and total predictive variance.
+
+        Inputs:
+            data (Tensor): Input data, shape [B, D]
+
+        Outputs:
+            If classification:
+                Tensor: Class probabilities, shape [B, num_classes]
+            If regression:
+                Tuple[Tensor, Tensor]: Output mean and total variance, shape [B, D_out]
+        """
+        
+        out_mean, out_var = self.forward_latent(data)
+
+        if self.likelihood == 'classification':
+            kappa = 1 / torch.sqrt(1. + np.pi / 8 * out_var)
+            return torch.softmax(kappa * out_mean, dim=-1)
+
+        if self.likelihood == 'regression':
+            return out_mean, out_var + self.sigma_noise ** 2
+    
+    def convert_model(self, org_model, posterior_variance):
+        """
+        Converts a deterministic MLP into a SUQ-compatible model with diagonal posterior.
+
+        Each layer is replaced with its corresponding SUQ module (e.g. linear, activation, batchnorm), using the provided flattened posterior variance vector.
+
+        Inputs:
+            org_model (nn.Module): The original model to convert (latent function only)
+            posterior_variance (Tensor): Flattened posterior variance for Bayesian parameters
+        """
+        
+        p_model = copy.deepcopy(org_model)
+
+        loc = 0
+        for n, layer in p_model.named_modules():
+            if isinstance(layer, nn.Linear):
+                
+                D_out, D_in = layer.weight.data.shape
+                num_param = torch.numel(parameters_to_vector(layer.parameters()))
+                num_weight_param = D_out * D_in
+                
+                covariance_block = posterior_variance[loc : loc + num_param]
+                
+                """
+                w_var: [D_out, D_in], w_var[k][i] = var(w_ki)
+                b_var: [D_out, ] b_var[k]: var(b_k)
+                """        
+                
+                b_var = torch.zeros_like(layer.bias.data).to(layer.bias.data.device)
+                w_var = torch.zeros_like(layer.weight.data).to(layer.bias.data.device)
+
+                for k in range(D_out):
+                    b_var[k] = covariance_block[num_weight_param + k]
+                    for i in range(D_in):
+                        w_var[k][i] = covariance_block[k * D_in + i]
+
+                new_layer = SUQ_Linear_Diag(layer, w_var, b_var)
+
+                loc += num_param
+                setattr(p_model, n, new_layer)
+            
+            if isinstance(layer, nn.BatchNorm1d):
+                new_layer = SUQ_BatchNorm_Diag(layer)
+                setattr(p_model, n, new_layer)
+                
+            if type(layer).__name__ in torch.nn.modules.activation.__all__:
+                new_layer = SUQ_Activation_Diag(layer)
+                setattr(p_model, n, new_layer)
+
+        self.model = p_model

suq/diag_suq_transformer.py

@@ -0,0 +1,627 @@
+import torch.nn as nn
+import torch.nn.functional as F
+import torch
+import numpy as np
+from suq.diag_suq_mlp import forward_aW_diag
+from suq.base_suq import SUQ_Base
+
+def forward_linear_diag_Bayesian_weight(e_mean, e_var, w_mean, w_var, bias = None):
+    """
+    Pass a distribution with diagonal covariance through a Bayesian linear layer with diagonal covariance
+
+    Given e ~ N(e_mean, e_cov), W ~ N(W_mean, W_cov), calculate the mean and variance h = eW.T + b.
+    
+    We only make the weight Bayesian, bias is treated determinstically
+
+    Note that as we always assume the input to next layer has diagonal covariance, so we only compute the variance over h here.
+    
+    Input
+        e_mean: [B, T, D_in] embedding mean
+        e_var: [B, T, D_in] embedding variance
+        w_mean: [D_out, D_in] weight mean
+        w_var: [D_out, D_in] weight covariance, w_cov[k][i]: var(w_ki)
+    Output
+        h_mean: [B, T, D_out]
+        h_var: [B, T, D_out] h_var[k] = var(h_k)
+    """
+
+    # calculate mean(h)
+    h_mean = F.linear(e_mean, w_mean, bias)
+    
+    # calculate var(h)
+    weight_mean2_var_sum = w_mean ** 2 + w_var # [D_out, D_in]
+    h_var = e_mean **2 @ w_var.T + e_var @ weight_mean2_var_sum.T
+
+    return h_mean, h_var
+
+def forward_linear_diag_determinstic_weight(e_mean, e_var, weight, bias = None):
+    """
+    Pass a distribution with diagonal covariance through a linear layer
+
+    Given e ~ N(e_mean, e_var) and determinstic weight W and bias b, calculate the mean and variance h = eW.T + b.
+    
+    Note that as we always assume the input to next layer has diagonal covariance, so we only compute the variance over h here.
+    
+    Input
+        e_mean: [B, T, D_in] embedding mean
+        e_var: [B, T, D_in] embedding variance
+        w_mean: [D_out, D_in] weight 
+    Output
+        h_mean: [B, T, D_out]
+        h_var: [B, T, D_out] h_var[k] = var(h_k)
+    """
+    
+    h_mean = F.linear(e_mean, weight, bias)
+    h_var = F.linear(e_var, weight ** 2, None)
+    
+    return h_mean, h_var
+
+@torch.enable_grad()
+def forward_activation_diag(activation_func, h_mean, h_var):
+    """
+    Pass a distribution with diagonal covariance through an activation layer. 
+
+    Given h ~ N(h_mean, h_cov), g(·), where h_cov is a diagonal matrix,
+    approximate the distribution of a = g(h) as 
+    a ~ N(g(h_mean), g'(h_mean)^T h_var g'(h_mean))
+    
+    Input
+        activation_func: g(·)
+        h_mean: [B, T, D] input mean
+        h_var: [B, T, D] input variance
+    
+    Output
+        a_mean: [B, T, D]
+        a_var: [B, T, D]
+    """
+
+    h_mean_grad = h_mean.detach().clone().requires_grad_()
+    
+    a_mean = activation_func(h_mean_grad)
+    a_mean.retain_grad()
+    a_mean.backward(torch.ones_like(a_mean)) #[B, T, D]
+    
+    nabla = h_mean_grad.grad #[B, T, D]
+    a_var = nabla ** 2 * h_var
+    
+    return a_mean.detach(), a_var
+    
+def forward_layer_norm_diag(e_mean, e_var, ln_weight, ln_eps):
+    """
+    Pass a distribution with diagonal covariance through LayerNorm layer
+    
+    Input
+        e_mean: mean of input distribution [B, T, D]
+        e_var: variance of input distribution [B, T, D]
+        ln_weight: layer norm scale factor
+        ln_eps: layer norm eps
+    
+    Output
+        output_var [B, T, D]
+    """
+
+    # calculate the var
+    input_mean_var = e_mean.var(dim=-1, keepdim=True, unbiased=False) # [B, T, 1]
+    scale_factor = (1 / (input_mean_var + ln_eps)) * ln_weight **2 # [B, T, D]
+    output_var = scale_factor * e_var # [B, T, D]
+    
+    return output_var
+
+def forward_value_cov_Bayesian_W(W_v, W_v_var, input_mean, input_var, n_h, D_v, diag_cov = False):
+    """
+    Given value matrix W_v ~ N(mean(W), var(W)) and input E ~ N(mean(E), var(E))
+    Compute the covariance of output v = W_v @ E 
+
+    Input: 
+        n_h: number of attention heads
+        D_v: dimension of value, n_h * D_v = D
+        W_v: value weight matrix [D, D]
+        W_v_var: variance of value matrix, [D, D]
+        input_mean: mean of input [B, T, D]
+        input_var: variance of input variance [B, T, D]
+        diag_cov: whether input only has diag covariance
+        
+    Output:
+        v_cov [B, T, n_h, D_v, D_v] or v_var [B, T, n_h, D_v]
+    """
+
+    B, T, D = input_var.size()
+
+    if not diag_cov:
+        ## compute general covariance 
+        W_v_reshaped = W_v.reshape(1, 1, n_h, D_v, D) 
+            # [D, D] -> [1, 1, n_h, D_v, D]
+        input_var_reshaped = input_var.reshape(B, T, 1, 1, D)
+            # [B, T, D] -> [B, T, 1, 1, D]
+        v_cov = (W_v_reshaped * input_var_reshaped).transpose(3, 4)
+            # [1, 1, n_h, D_v, D] * [B, T, 1, 1, D] -> [B, T, n_h, D_v, D] -> [B, T, n_h, D, D_v]
+        v_cov = torch.matmul(W_v_reshaped, v_cov)
+            #  [1, 1, n_h, D_v, D] @ [B, T, n_h, D, D_v]  -> [B, T, n_h, D_v, D_v]
+        
+        ## add missing part for variance
+        W_v_var_reshaped = W_v_var.reshape(1, 1, n_h, D_v, D) 
+            #[D, D] -> [1, 1, n_h, D_v, D]
+        input_var_plus_mean_square = input_var_reshaped + input_mean.reshape(B, T, 1, 1, D)**2 #[B, T, 1, 1, D]
+        extra_var_term = torch.sum(input_var_plus_mean_square * W_v_var_reshaped, dim=[4]) # [B, T, n_h, D_v, D] -> [B, T, n_h, D_v]
+        v_cov = v_cov + torch.diag_embed(extra_var_term) 
+
+        return v_cov
+
+    else:
+        weight_mean2_var_sum = W_v **2 + W_v_var # [D, D]
+        v_var = input_mean **2 @ W_v_var.T + input_var @ weight_mean2_var_sum.T # [B, T, D]
+    
+        return v_var.reshape(B, T, n_h, D_v)
+
+def forward_value_cov_determinstic_W(W_v, input_var, n_h, D_v, diag_cov = False):
+    """
+    Given determinstic value matrix W_v and input E ~ N(mean(E), var(E))
+    Compute the covariance of output v = W_v @ E 
+
+    
+    Input: 
+        n_h: number of attention heads
+        D_v: dimension of value, n_h * D_v = D
+        W_v: value weight matrix [D, D], which can be reshaped into [n_h, D_v, D]
+        input_var: variance of input variance [B, T, D]
+        diag_cov: whether input only has diag covariance
+        
+    Output:
+        v_cov [B, T, n_h, D_v, D_v] or v_var [B, T, n_h, D_v]
+    """
+
+    B, T, D = input_var.size()
+
+    if not diag_cov:
+        W_v_reshaped = W_v.reshape(1, 1, n_h, D_v, D) 
+            #[n_h, D_v, D] -> [1, 1, n_h, D_v, D]
+        input_var_reshaped = input_var.reshape(B, T, 1, 1, D)
+            # [B, T, D] -> [B, T, 1, 1, D]
+        v_cov = (W_v_reshaped * input_var_reshaped).transpose(3, 4)
+            # [1, 1, n_h, D_v, D] * [B, T, 1, 1, D] -> [B, T, n_h, D_v, D] -> [B, T, n_h, D, D_v]
+        v_cov = torch.matmul(W_v_reshaped, v_cov)
+            #  [1, 1, n_h, D_v, D] @ [B, T, n_h, D, D_v]  -> [B, T, n_h, D_v, D_v]
+
+        return v_cov
+    
+    else:
+        v_var = input_var @ (W_v ** 2).T
+
+        return v_var.reshape(B, T, n_h, D_v)
+    
+def forward_QKV_cov(attention_score, v_cov, diag_cov = False):
+    """
+    given attention score (QK^T) and V ~ N(mean(V), cov(V))
+    compute the covariance of output E = (QK^T) V
+    
+    Input:
+        attention_score: [B, n_h, T, T] attention_score[t] is token t's attention score for all other tokens
+        v_cov: [B, T, n_h, D_v, D_v] or [B, T, n_h, D_v], covariance of value
+        diag_cov: whether input only has diag covariance
+        
+    Output:
+        QKV_cov: [B, n_h, T, D_v, D_v] or [B, T, n_h, D_v], covariance of output E
+    """
+    if diag_cov:
+        B, T, n_h, D_v = v_cov.size()
+        QKV_cov = attention_score **2 @ v_cov.transpose(1, 2) # [B, n_h, T, D_v]
+            # v_cov [B, T, n_h, D_v] -> [B, n_h, T, D_v]
+            # [B, n_h, T, T] @ [B, n_h, T, D_v]  -> [B, n_h, T, D_v]
+    else:
+        
+        B, T, n_h, D_v, _ = v_cov.size()
+        
+        QKV_cov = attention_score **2 @ v_cov.permute(0, 2, 1, 3, 4).reshape(B, n_h, T, D_v * D_v) # [B, n_h, T, D_v * D_v]
+        # v_cov [B, T, n_h, D_v, D_v] -> [B, n_h, T, D_v * D_v]
+        # [B, n_h, T, T] @ [B, n_h, T, D_v * D_v]  -> [B, n_h, T, D_v * D_v]
+        QKV_cov = QKV_cov.reshape(B, n_h, T, D_v, D_v)
+        
+    return QKV_cov
+
+def forward_fuse_multi_head_cov(QKV_cov, project_W, diag_cov = False):
+    """
+    given concatanated multi-head embedding E ~ N(mean(E), cov(E)) and project weight matrix W
+    compute variance of each output dimenison
+    
+    Input:
+        QKV_cov: [B, n_h, T, D_v, D_v] or [B, n_h, T, D_v]
+        project_W: [D, D]  D_out x D_in (n_h * D_v)
+        diag_cov: whether input only has diag covariance
+        
+    Output: 
+        output_var [B, T, D]
+    """
+    if diag_cov:
+        B, n_h, T, D_v = QKV_cov.size()
+        output_var = QKV_cov.permute(0, 2, 1, 3).reshape(B, T, n_h * D_v) @ project_W ** 2
+            # QKV_cov [B, n_h, T, D_v] -> [B, T, n_h, D_v] -> [B, T, n_h * D_v]
+
+        return output_var
+        
+    else:
+        B, n_h, T, D_v, _ = QKV_cov.size()
+        D, _ = project_W.shape
+        
+        project_W_reshaped_1 = project_W.T.reshape(n_h, D_v, D).permute(0, 2, 1).reshape(n_h * D, D_v, 1)
+            # [n_h, D_v, D] -> [n_h, D, D_v] -> [n_h * D, D_v, 1]
+        project_W_reshaped_2 = project_W.T.reshape(n_h, D_v, D).permute(0, 2, 1).reshape(n_h * D, 1, D_v)
+            # [n_h, D_v, D] -> [n_h, D, D_v] -> [n_h * D, 1, D_v]
+
+        project_W_outer = torch.bmm(project_W_reshaped_1, project_W_reshaped_2).reshape(n_h, D, D_v, D_v).permute(1, 0, 2, 3) # [D, n_h, D_v, D_v]
+        # [n_h * D, D_v, D_v] @ [n_h * D, 1, D_v] -> [n_h * D, D_v, D_v] -> [D, n_h, D_v, D_v]
+            
+        output_var_einsum = torch.einsum('dhij,bthij->dbt', project_W_outer, QKV_cov.permute(0, 2, 1, 3, 4))
+        
+        return output_var_einsum.permute(1, 2, 0)
+
+class SUQ_LayerNorm_Diag(nn.Module):
+    """
+    LayerNorm module with uncertainty propagation under SUQ.
+
+    Wraps `nn.LayerNorm` and propagates input variance analytically using running statistics. See the SUQ paper for theoretical background and assumptions.
+
+    Inputs:
+        LayerNorm (nn.LayerNorm): The original layer norm module to wrap
+    """
+
+    def __init__(self, LayerNorm):
+        super().__init__()
+        
+        self.LayerNorm = LayerNorm
+    
+    def forward(self, x_mean, x_var):
+        """
+        Inputs:
+            x_mean (Tensor): Input mean, shape [B, T, D]
+            x_var (Tensor): Input variance, shape [B, T, D]
+
+        Outputs:
+            out_mean (Tensor): Output mean after layer norm, shape [B, T, D]
+            out_var (Tensor): Output variance after layer norm, shape [B, T, D]
+        """
+        with torch.no_grad():
+        
+            out_mean = self.LayerNorm.forward(x_mean)
+            out_var = forward_layer_norm_diag(x_mean, x_var, self.LayerNorm.weight, 1e-5)
+            
+        return out_mean, out_var
+
+
+class SUQ_Classifier_Diag(nn.Module):
+    """
+    Classifier head with uncertainty propagation under SUQ, with a diagonal Gaussian posterior.
+
+    Wraps a standard linear classifier and applies closed-form mean and variance propagation.
+    See the SUQ paper for theoretical background and assumptions.
+
+    Inputs:
+        classifier (nn.Linear): The final classification head
+        w_var (Tensor): Weight variances, shape [D_out, D_in]
+        b_var (Tensor): Bias variances, shape [D_out]
+    """
+
+    def __init__(self, classifier, w_var, b_var):
+        super().__init__()
+        
+        self.weight = classifier.weight
+        self.bias = classifier.bias
+        self.w_var = w_var.reshape(self.weight.shape)
+        self.b_var = b_var.reshape(self.bias.shape)
+    
+    def forward(self, x_mean, x_var):
+        """
+        Inputs:
+            x_mean (Tensor): Input mean, shape [B, D]
+            x_var (Tensor): Input variance, shape [B, D]
+
+        Outputs:
+            h_mean (Tensor): Output mean, shape [B, D_out]
+            h_var (Tensor): Output variance, shape [B, D_out]
+        """
+        with torch.no_grad():
+            h_mean, h_var = forward_aW_diag(x_mean, x_var, self.weight.data, self.bias.data, self.w_var, self.b_var)
+        return h_mean, h_var
+
+class SUQ_TransformerMLP_Diag(nn.Module):
+    """
+    MLP submodule of a transformer block with uncertainty propagation under SUQ.
+
+    Supports both deterministic and Bayesian forward modes with closed-form variance propagation.
+    Used internally in `SUQ_Transformer_Block_Diag`.
+
+    Inputs:
+        MLP (nn.Module): Original MLP submodule
+        determinstic (bool): Whether to treat the MLP weights as deterministic
+        w_fc_var (Tensor, optional): Variance of the first linear layer (if Bayesian)
+        w_proj_var (Tensor, optional): Variance of the second linear layer (if Bayesian)
+    """
+
+    def __init__(self, MLP, determinstic = True, w_fc_var = None, w_proj_var = None):
+        super().__init__()
+
+        self.MLP = MLP
+        self.determinstic = determinstic
+        if not determinstic:
+            self.w_fc_var = w_fc_var.reshape(self.MLP.c_fc.weight.shape)
+            self.w_proj_var = w_proj_var.reshape(self.MLP.c_proj.weight.shape)
+
+    def forward(self, x_mean, x_var):
+        """
+        Inputs:
+            x_mean (Tensor): Input mean, shape [B, T, D]
+            x_var (Tensor): Input variance, shape [B, T, D]
+
+        Outputs:
+            h_mean (Tensor): Output mean, shape [B, T, D]
+            h_var (Tensor): Output variance, shape [B, T, D]
+        """
+        
+        # first fc layer
+        with torch.no_grad():
+            if self.determinstic:
+                h_mean, h_var = forward_linear_diag_determinstic_weight(x_mean, x_var, self.MLP.c_fc.weight.data, self.MLP.c_fc.bias.data)
+            else:
+                h_mean, h_var = forward_linear_diag_Bayesian_weight(x_mean, x_var, self.MLP.c_fc.weight.data, self.w_fc_var, self.MLP.c_fc.bias.data)
+        # activation function
+        h_mean, h_var = forward_activation_diag(self.MLP.gelu, h_mean, h_var)
+        # second fc layer
+        with torch.no_grad():
+            if self.determinstic:
+                h_mean, h_var = forward_linear_diag_determinstic_weight(h_mean, h_var, self.MLP.c_proj.weight.data, self.MLP.c_proj.bias.data)
+            else:
+                h_mean, h_var = forward_linear_diag_Bayesian_weight(h_mean, h_var, self.MLP.c_proj.weight.data, self.w_proj_var, self.MLP.c_proj.bias.data)
+
+        return h_mean, h_var
+
+class SUQ_Attention_Diag(nn.Module):
+    """
+    Self-attention module with uncertainty propagation under SUQ.
+
+    Supports deterministic and Bayesian value projections, with optional diagonal covariance assumptions. For details see SUQ paper section A.6
+    Used internally in `SUQ_Transformer_Block_Diag`.
+
+    Inputs:
+        Attention (nn.Module): The original attention module
+        determinstic (bool): Whether to treat value projections as deterministic
+        diag_cov (bool): If True, only compute the diagoanl covariance for value
+        W_v_var (Tensor, optional): Posterior variance for value matrix (if Bayesian)
+    """
+    
+    def __init__(self, Attention, determinstic = True, diag_cov = False, W_v_var = None):
+        super().__init__()
+        
+        self.Attention = Attention
+        self.determinstic = determinstic
+        self.diag_cov = diag_cov
+        
+        if not self.determinstic:
+            self.W_v_var = W_v_var # [D * D]
+
+    def forward(self, x_mean, x_var):
+        """
+        Inputs:
+            x_mean (Tensor): Input mean, shape [B, T, D]
+            x_var (Tensor): Input variance, shape [B, T, D]
+
+        Outputs:
+            output_mean (Tensor): Output mean after attention, shape [B, T, D]
+            output_var (Tensor): Output variance after attention, shape [B, T, D]
+        """
+
+        with torch.no_grad():
+        
+            output_mean, attention_score = self.Attention.forward(x_mean, True)
+            
+            n_h = self.Attention.n_head
+            B, T, D = x_mean.size()
+            D_v = D // n_h
+            
+            W_v = self.Attention.c_attn_v.weight.data
+            project_W = self.Attention.c_proj.weight.data
+
+            if self.determinstic:
+                v_cov = forward_value_cov_determinstic_W(W_v, x_var, n_h, D_v)
+            else:
+                v_cov = forward_value_cov_Bayesian_W(W_v, self.W_v_var.reshape(D, D), x_mean, x_var, n_h, D_v, self.diag_cov)
+
+            QKV_cov = forward_QKV_cov(attention_score, v_cov, self.diag_cov)
+            output_var = forward_fuse_multi_head_cov(QKV_cov, project_W, self.diag_cov)
+
+            return output_mean, output_var
+        
+class SUQ_Transformer_Block_Diag(nn.Module):
+    """
+    Single transformer block with uncertainty propagation under SUQ.
+
+    Wraps LayerNorm, attention, and MLP submodules with uncertainty-aware versions.
+    Used in `SUQ_ViT_Diag` to form a full transformer stack.
+
+    Inputs:
+        MLP (nn.Module): Original MLP submodule
+        Attention (nn.Module): Original attention submodule
+        LN_1 (nn.LayerNorm): Pre-attention layer norm
+        LN_2 (nn.LayerNorm): Pre-MLP layer norm
+        MLP_determinstic (bool): Whether to treat MLP as deterministic
+        Attn_determinstic (bool): Whether to treat attention as deterministic
+        diag_cov (bool): If True, only compute the diagoanl covariance for value
+        w_fc_var (Tensor or None): Posterior variance of MLP input projection (if Bayesian)
+        w_proj_var (Tensor or None): Posterior variance of MLP output projection (if Bayesian)
+        W_v_var (Tensor or None): Posterior variance of value matrix (if Bayesian)
+    """
+
+    
+    def __init__(self, MLP, Attention, LN_1, LN_2, MLP_determinstic, Attn_determinstic, diag_cov = False, w_fc_var = None, w_proj_var = None, W_v_var = None):
+        super().__init__()
+        
+        self.ln_1 = SUQ_LayerNorm_Diag(LN_1)
+        self.ln_2 = SUQ_LayerNorm_Diag(LN_2)
+        self.attn = SUQ_Attention_Diag(Attention, Attn_determinstic, diag_cov, W_v_var)
+        self.mlp = SUQ_TransformerMLP_Diag(MLP, MLP_determinstic, w_fc_var, w_proj_var)
+
+    def forward(self, x_mean, x_var):
+        """
+        Inputs:
+            x_mean (Tensor): Input mean, shape [B, T, D]
+            x_var (Tensor): Input variance, shape [B, T, D]
+
+        Outputs:
+            h_mean (Tensor): Output mean after transformer block, shape [B, T, D]
+            h_var (Tensor): Output variance after transformer block, shape [B, T, D]
+        """
+        
+        h_mean, h_var = self.ln_1(x_mean, x_var)
+        h_mean, h_var = self.attn(h_mean, h_var)
+        h_mean = h_mean + x_mean
+        h_var = h_var + x_var
+        
+        old_h_mean, old_h_var = h_mean, h_var
+        
+        h_mean, h_var = self.ln_2(h_mean, h_var)
+        h_mean, h_var = self.mlp(h_mean, h_var)
+        h_mean = h_mean + old_h_mean
+        h_var = h_var + old_h_var
+        
+        return h_mean, h_var
+
+
+class SUQ_ViT_Diag(SUQ_Base):
+    """
+    Vision Transformer model with uncertainty propagation under SUQ, with a diagonal Gaussian posterior.
+
+    Wraps a ViT architecture into a structured uncertainty-aware model by replacing parts
+    of the network with SUQ-compatible blocks. Allows selective Bayesian treatment of MLP
+    and attention modules within each transformer block.
+
+    Currently supports classification only. See the SUQ paper for theoretical background and assumptions.
+
+    Inputs:
+        ViT (nn.Module): A Vision Transformer model structured like `examples/vit_model.py`
+        posterior_variance (Tensor): Flattened posterior variance vector
+        MLP_determinstic (bool): Whether MLP submodules are treated as deterministic
+        Attn_determinstic (bool): Whether attention submodules are treated as deterministic
+        scale_init (float, optional): Initial value for the scale factor
+        attention_diag_cov (bool): If True, only compute the diagoanl covariance for value
+        likelihood (str): Currently only support 'Classification'
+        num_det_blocks (int): Number of transformer blocks to leave deterministic (from the bottom up)
+    """
+
+    def __init__(self, ViT, posterior_variance, MLP_determinstic, Attn_determinstic, scale_init = 1.0, attention_diag_cov = False, likelihood = 'clasification', num_det_blocks = 10):
+        super().__init__(likelihood, scale_init)
+        
+        if likelihood not in ['classification']:
+            raise ValueError(f"{likelihood} not supported for ViT")
+        
+    
+        self.transformer = nn.ModuleDict(dict(
+            pte = ViT.transformer.pte,
+            h = nn.ModuleList(),
+            ln_f = SUQ_LayerNorm_Diag(ViT.transformer.ln_f)
+        ))
+        
+        self.scale_factor = nn.Parameter(torch.Tensor([scale_init]))
+        
+        num_param_c_fc = ViT.transformer.h[0].mlp.c_fc.weight.numel()
+        num_param_c_proj = ViT.transformer.h[0].mlp.c_proj.weight.numel()
+        num_param_value_matrix = ViT.transformer.h[0].attn.c_proj.weight.numel()
+        
+        index = 0
+        for block_index in range(len(ViT.transformer.h)):
+            
+            if block_index < num_det_blocks:
+                self.transformer.h.append(ViT.transformer.h[block_index])
+            else:
+                if not MLP_determinstic:
+                    w_fc_var = posterior_variance[index: index + num_param_c_fc]
+                    index += num_param_c_fc
+                    w_proj_var = posterior_variance[index: index + num_param_c_proj]
+                    index += num_param_c_proj
+                    self.transformer.h.append(
+                        SUQ_Transformer_Block_Diag(ViT.transformer.h[block_index].mlp, 
+                                                    ViT.transformer.h[block_index].attn, 
+                                                    ViT.transformer.h[block_index].ln_1, 
+                                                    ViT.transformer.h[block_index].ln_2, 
+                                                    MLP_determinstic,
+                                                    Attn_determinstic,
+                                                    attention_diag_cov,
+                                                    w_fc_var, 
+                                                    w_proj_var,
+                                                    None))
+                
+                if not Attn_determinstic:
+                    w_v_var = posterior_variance[index : index + num_param_value_matrix]
+                    index += num_param_value_matrix
+                    self.transformer.h.append(
+                        SUQ_Transformer_Block_Diag(ViT.transformer.h[block_index].mlp, 
+                                                    ViT.transformer.h[block_index].attn, 
+                                                    ViT.transformer.h[block_index].ln_1, 
+                                                    ViT.transformer.h[block_index].ln_2, 
+                                                    MLP_determinstic,
+                                                    Attn_determinstic,
+                                                    attention_diag_cov,
+                                                    None, 
+                                                    None,
+                                                    w_v_var))
+
+        num_param_classifier_weight = ViT.classifier.weight.numel()
+        self.classifier = SUQ_Classifier_Diag(ViT.classifier, posterior_variance[index: index + num_param_classifier_weight], posterior_variance[index + num_param_classifier_weight:])
+
+    def forward_latent(self, pixel_values, interpolate_pos_encoding = None):
+        
+        """
+        Compute the predictive mean and variance of the ViT's latent output before applying the final likelihood layer.
+
+        Traverses the full transformer stack with uncertainty propagation.
+
+        Inputs:
+            pixel_values (Tensor): Input image tensor, shape [B, C, H, W]
+            interpolate_pos_encoding (optional): Optional positional embedding interpolation
+
+        Outputs:
+            x_mean (Tensor): Predicted latent mean at the [CLS] token, shape [B, D]
+            x_var (Tensor): Predicted latent variance at the [CLS] token, shape [B, D]
+        """
+
+        device = pixel_values.device
+
+        x_mean = self.transformer.pte(
+            pixel_values, interpolate_pos_encoding=interpolate_pos_encoding
+        )
+        
+        # pass through model
+        x_var = torch.zeros_like(x_mean, device = device)
+        
+        for i, block in enumerate(self.transformer.h):
+            
+            if isinstance(block, SUQ_Transformer_Block_Diag):
+            
+                x_mean, x_var = block(x_mean, x_var)
+            else:
+                x_mean = block(x_mean)
+
+        x_mean, x_var = self.transformer.ln_f(x_mean, x_var)
+        
+        x_mean, x_var = self.classifier(x_mean[:, 0, :], x_var[:, 0, :])
+        x_var = x_var / self.scale_factor.to(device)
+        
+        return x_mean, x_var
+
+    def forward(self, pixel_values, interpolate_pos_encoding = None):
+        """
+        Compute predictive class probabilities using a probit approximation.
+
+        Performs a full forward pass through the ViT with uncertainty propagation, and
+        produces softmax-normalized class probabilities for classification.
+
+        Inputs:
+            pixel_values (Tensor): Input image tensor, shape [B, C, H, W]
+            interpolate_pos_encoding (optional): Optional positional embedding interpolation
+
+        Outputs:
+            Tensor: Predicted class probabilities, shape [B, num_classes]
+        """
+
+        x_mean, x_var = self.forward_latent(pixel_values, interpolate_pos_encoding)
+        kappa = 1 / torch.sqrt(1. + np.pi / 8 * x_var)
+        
+        return torch.softmax(kappa * x_mean, dim=-1)
+

suq/streamline_layer.py

@@ -0,0 +1,23 @@
+from .diag_suq_mlp import (
+    SUQ_Linear_Diag,
+    SUQ_Activation_Diag,
+    SUQ_BatchNorm_Diag
+)
+from .diag_suq_transformer import (
+    SUQ_TransformerMLP_Diag,
+    SUQ_Attention_Diag,
+    SUQ_LayerNorm_Diag,
+    SUQ_Classifier_Diag,
+    SUQ_Transformer_Block_Diag
+)
+
+__all__ = [
+    "SUQ_Linear_Diag",
+    "SUQ_Activation_Diag",
+    "SUQ_BatchNorm_Diag",
+    "SUQ_TransformerMLP_Diag",
+    "SUQ_Attention_Diag",
+    "SUQ_LayerNorm_Diag",
+    "SUQ_Classifier_Diag",
+    "SUQ_Transformer_Block_Diag"
+]

suq-0.1.0.dist-info/METADATA

@@ -0,0 +1,19 @@
+Metadata-Version: 2.4
+Name: suq
+Version: 0.1.0
+Summary: Streamlined Uncertainty Quantification (SUQ)
+Home-page: https://github.com/AaltoML/SUQ
+Author: Rui Li, Marcus Klasson, Arno Solin, Martin Trapp
+License: MIT
+Classifier: Programming Language :: Python :: 3
+License-File: LICENSE
+Requires-Dist: torch>=1.10
+Requires-Dist: numpy>=1.21
+Requires-Dist: tqdm>=4.60
+Dynamic: author
+Dynamic: classifier
+Dynamic: home-page
+Dynamic: license
+Dynamic: license-file
+Dynamic: requires-dist
+Dynamic: summary

suq-0.1.0.dist-info/RECORD

@@ -0,0 +1,12 @@
+suq/SUQ_MLP.py,sha256=ciegIOz2Y0vtJ4Uc56dhwHjeaAoLAKqQQnSTgyh8Sqc,497
+suq/SUQ_ViT.py,sha256=6BpHMOLf1qhVRgqM3guAcS60PpYJ1L1CEWdIerIKrv4,841
+suq/__init__.py,sha256=5LGGZQ6wwjEfwsmzx55lSA4nfBx1N1qQpzgFff8zpJM,120
+suq/base_suq.py,sha256=77LggEGn5m3h472kRLVghM1j7WCY4zJih3undVdIoS0,7963
+suq/diag_suq_mlp.py,sha256=yE4p2pQtD97g6GvebkROqNDPFuJUhrwZH9zRkWAiSWk,10985
+suq/diag_suq_transformer.py,sha256=QfGkPCf1DJ1cVIrkDp7tD8bImtCgt3XoUnngtgfwIAc,25576
+suq/streamline_layer.py,sha256=sew0BodioNjElVpSq0c87c0ORRA5yGDbLTy1eeLp8VI,504
+suq-0.1.0.dist-info/licenses/LICENSE,sha256=XbdRUcHQPFVqh9MgS85KpgRmx09FAPBW5gOHyzjJz3c,1064
+suq-0.1.0.dist-info/METADATA,sha256=lcWlRUIWAM-LKK8MqbDehGjJ3QX1_Fl7nDCtZ-Jb1V4,497
+suq-0.1.0.dist-info/WHEEL,sha256=CmyFI0kx5cdEMTLiONQRbGQwjIoR1aIYB7eCAQ4KPJ0,91
+suq-0.1.0.dist-info/top_level.txt,sha256=_8edEeGJ6W1u-AhLS6x16sRWGGKwumOGPl7EBsW-v5g,4
+suq-0.1.0.dist-info/RECORD,,

suq-0.1.0.dist-info/WHEEL

@@ -0,0 +1,5 @@
+Wheel-Version: 1.0
+Generator: setuptools (78.1.0)
+Root-Is-Purelib: true
+Tag: py3-none-any
+

suq-0.1.0.dist-info/licenses/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2025 AaltoML
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

suq-0.1.0.dist-info/top_level.txt

@@ -0,0 +1 @@
+suq