BayesForge 0.0.46__tar.gz

bayesforge-0.0.46/BF/BNN/activations.py

@@ -0,0 +1,249 @@
+import jax
+import jax.numpy as jnp
+from jax.scipy.special import erf
+
+class activation:
+    """
+    A collection of common neural network activation functions implemented in JAX.
+
+    This class serves as a namespace for activation functions, which can be
+    called directly as static methods (e.g., `activation.relu(x)`).
+    """
+
+    # --- Original Functions (with improved docstrings) ---
+
+    @staticmethod
+    def tanh(x: jnp.ndarray) -> jnp.ndarray:
+        """
+        Computes the hyperbolic tangent activation function.
+
+        This function squashes its input values into the range [-1, 1]. It is
+        zero-centered, which can help with optimization by making the mean of
+        the activations closer to zero. However, it can suffer from the
+        vanishing gradient problem for very large or very small inputs.
+
+        Parameters:
+        - x (jnp.ndarray): The input array.
+
+        Returns:
+        - jnp.ndarray: The array with tanh applied element-wise.
+        """
+        return jnp.tanh(x)
+
+    @staticmethod
+    def relu(x: jnp.ndarray) -> jnp.ndarray:
+        """
+        Computes the Rectified Linear Unit (ReLU) activation function.
+
+        ReLU is one of the most widely used activation functions. It computes
+        f(x) = max(0, x), effectively "turning off" neurons with negative
+        outputs. This makes the network sparse and computationally efficient.
+        Its main drawback is the "dying ReLU" problem, where neurons can get
+        stuck in a state where they always output zero if their weights are
+        updated such that their input is always negative.
+
+        Parameters:
+        - x (jnp.ndarray): The input array.
+
+        Returns:
+        - jnp.ndarray: The array with ReLU applied element-wise.
+        """
+        return jnp.maximum(0, x)
+
+    @staticmethod
+    def sigmoid(x: jnp.ndarray) -> jnp.ndarray:
+        """
+        Computes the sigmoid (or logistic) activation function.
+
+        This function squashes its input values into the range [0, 1]. It is
+        commonly used in the output layer of binary classification models to
+        represent a probability. It is less favored for hidden layers due to
+        its non-zero-centered output and strong saturation, which leads to
+        vanishing gradients.
+
+        Parameters:
+        - x (jnp.ndarray): The input array.
+
+        Returns:
+        - jnp.ndarray: The array with sigmoid applied element-wise.
+        """
+        return jax.nn.sigmoid(x)
+
+    @staticmethod
+    def softmax(x: jnp.ndarray, axis: int = -1) -> jnp.ndarray:
+        """
+        Computes the softmax activation function.
+
+        Softmax transforms a vector of real numbers into a probability
+        distribution. Each output value is in the range [0, 1], and all
+        output values sum to 1. It is almost exclusively used as the output
+        activation function for multi-class classification problems.
+
+        Parameters:
+        - x (jnp.ndarray): The input array.
+        - axis (int): The axis along which the softmax should be computed.
+
+        Returns:
+        - jnp.ndarray: The array with softmax applied along the specified axis.
+        """
+        return jax.nn.softmax(x, axis=axis)
+
+    # --- New and Important Activation Functions ---
+
+    @staticmethod
+    def leaky_relu(x: jnp.ndarray, negative_slope: float = 0.01) -> jnp.ndarray:
+        """
+        Computes the Leaky Rectified Linear Unit (Leaky ReLU).
+
+        Leaky ReLU is a variant of ReLU designed to solve the "dying ReLU"
+        problem. Instead of being zero for negative inputs, it allows a small,
+        non-zero gradient (controlled by `negative_slope`). This ensures that
+        neurons do not become completely inactive.
+
+        f(x) = x if x > 0, else negative_slope * x.
+
+        Usage:
+        A common drop-in replacement for ReLU, especially if you suspect
+        dying neurons are an issue.
+
+        Parameters:
+        - x (jnp.ndarray): The input array.
+        - negative_slope (float): The small slope for negative inputs. Default is 0.01.
+
+        Returns:
+        - jnp.ndarray: The array with Leaky ReLU applied element-wise.
+        """
+        return jnp.where(x >= 0, x, negative_slope * x)
+
+    @staticmethod
+    def elu(x: jnp.ndarray, alpha: float = 1.0) -> jnp.ndarray:
+        """
+        Computes the Exponential Linear Unit (ELU).
+
+        ELU is another alternative to ReLU that also aims to solve the dying
+        neuron problem and can lead to faster learning. For negative inputs,
+        it becomes a smooth, saturating function that pushes the mean
+        activation closer to zero, which can speed up convergence.
+
+        f(x) = x if x > 0, else alpha * (exp(x) - 1).
+
+        Usage:
+        Often provides better performance than ReLU or Leaky ReLU but is
+        slightly more computationally expensive due to the exponential function.
+
+        Parameters:
+        - x (jnp.ndarray): The input array.
+        - alpha (float): The saturation parameter for negative inputs. Default is 1.0.
+
+        Returns:
+        - jnp.ndarray: The array with ELU applied element-wise.
+        """
+        return jnp.where(x >= 0, x, alpha * (jnp.exp(x) - 1))
+
+    @staticmethod
+    def gelu(x: jnp.ndarray) -> jnp.ndarray:
+        """
+        Computes the Gaussian Error Linear Unit (GELU).
+
+        GELU is a smooth, high-performing activation function that is the
+        standard in modern Transformer models like BERT and GPT. It weights
+        its input by its value, but this weighting is stochastic and depends
+        on the standard Gaussian cumulative distribution function (CDF).
+        Intuitively, it's more likely to "drop" (zero-out) inputs that are
+        closer to zero.
+
+        f(x) = x * Φ(x), where Φ(x) is the standard normal CDF.
+
+        Usage:
+        The state-of-the-art choice for Transformer-based architectures. A
+        strong general-purpose choice for many deep networks.
+
+        Parameters:
+        - x (jnp.ndarray): The input array.
+
+        Returns:
+        - jnp.ndarray: The array with GELU applied element-wise.
+        """
+        return x * 0.5 * (1.0 + erf(x / jnp.sqrt(2.0)))
+
+    @staticmethod
+    def silu(x: jnp.ndarray) -> jnp.ndarray:
+        """
+        Computes the Sigmoid Linear Unit (SiLU), also known as Swish.
+
+        SiLU is a self-gated activation function, defined as f(x) = x * sigmoid(x).
+        The sigmoid part acts as a soft "gate" that modulates the input. This
+        function is smooth, non-monotonic, and has been shown to perform as
+        well or better than ReLU on many challenging tasks without adding any
+        extra parameters.
+
+        Usage:
+        An excellent, modern, general-purpose replacement for ReLU that often
+        improves performance.
+
+        Parameters:
+        - x (jnp.ndarray): The input array.
+
+        Returns:
+        - jnp.ndarray: The array with SiLU applied element-wise.
+        """
+        return x * jax.nn.sigmoid(x)
+
+    @staticmethod
+    def softplus(x: jnp.ndarray) -> jnp.ndarray:
+        """
+        Computes the Softplus activation function.
+
+        Softplus is a smooth approximation of the ReLU function, defined as
+        f(x) = log(1 + exp(x)). Its output is always strictly positive.
+
+        Usage:
+        While less common as a hidden layer activation, it is very useful in
+        the output layer of a model when a strictly positive output is
+        required, for example, when predicting the variance (scale) parameter
+        of a distribution.
+
+        Parameters:
+        - x (jnp.ndarray): The input array.
+
+        Returns:
+        - jnp.ndarray: The array with Softplus applied element-wise.
+        """
+        return jnp.log(1.0 + jnp.exp(x))
+        # A more numerically stable version is: jax.nn.softplus(x)
+
+    @staticmethod
+    def linear(x: jnp.ndarray) -> jnp.ndarray:
+        """
+        Computes the linear (or identity) activation function.
+
+        This function simply returns the input without any modification, i.e.,
+        f(x) = x.
+
+        Usage:
+        This is a critical component. It is the default choice for the output
+        layer of any regression model, where the output is an unbounded,
+        continuous value.
+
+        Parameters:
+        - x (jnp.ndarray): The input array.
+
+        Returns:
+        - jnp.ndarray: The identical input array.
+        """
+        return x
+
+    @staticmethod
+    def exp(x: jnp.ndarray) -> jnp.ndarray:
+        """
+        Computes the exponential activation function.
+
+        f(x) = exp(x). Output is always positive.
+
+        Parameters:
+        - x (jnp.ndarray): The input array.
+
+        Returns:
+        - jnp.ndarray: The array with exponential applied element-wise.
+        """
+        return jnp.exp(x)

bayesforge-0.0.46/BF/BNN/bnn.py

@@ -0,0 +1,272 @@
+from BF.Distributions.np_dists import UnifiedDist as dist
+from BF.BNN.activations import activation
+
+from numpyro import deterministic
+import jax.numpy as jnp
+import jax
+
+class bnn(activation):
+    """
+    The bnn class is designed to build Bayesian Neural Networks (BNNs). It provides methods for creating network layers with specified prior distributions and activation functions. Additionally, it includes a specific two-layer BNN model for covariance estimation and a utility function to compute a correlation matrix from posterior samples.
+    """
+    def __init__(self, rand_seed = True):
+        super().__init__()
+
+        self.dist = dist(rand_seed)
+        # Create the mapping in the constructor
+        self._activation_map = {
+            # Standard functions
+            "relu": self.relu,
+            "tanh": self.tanh,
+            "sigmoid": self.sigmoid,
+            "softmax": self.softmax,
+            "linear": self.linear,
+
+            # Advanced/Modern functions
+            "leaky_relu": self.leaky_relu,
+            "elu": self.elu,
+            "gelu": self.gelu,
+            "silu": self.silu,
+            "swish": self.silu,  # Common alias for SiLU
+
+            # Specialty functions
+            "softplus": self.softplus,
+            "exp": self.exp,
+        }
+
+    def available_activations(self):
+        """
+        Returns a list of available activation functions.
+        
+        This method retrieves the names of all activation functions defined in the class.
+        """
+        return list(self._activation_map.keys())
+
+    def layer_linear(self, X, dist, activation=None, bias=False):
+        """        
+        Adds a layer to the BNN with the specified prior distribution and activation function.       
+
+        Parameters:
+        - prior_dist (BF.dist): The prior distribution for the weights of the layer. The shape of the distribution defines the layer's input/output dimensions.
+        - activation (str): The name of the activation function to use after this layer ('relu', 'tanh', 'sigmoid', 'softmax').  
+        """
+        print("⚠️This function is still in development. Use it with caution. ⚠️")
+        if bias is not False:
+            prod = jnp.matmul(X, dist) + bias
+        else:
+            prod = jnp.matmul(X, dist) 
+
+        # 2. Get and store the activation function object.
+        if activation is None:
+            return prod
+        else:
+            try:
+                activation_func = getattr(self, activation)
+            except AttributeError:
+                raise ValueError(f"Unknown activation function: '{activation}'")    
+            return activation_func(prod)
+
+    def layer_attention(self, b_kv, b_q, d_model=32,  sample=True, name = '', seed = None):
+        print("⚠️This function is still in development. Use it with caution. ⚠️")
+        # Layers
+        ### Dimensions 
+        self.b_q, self.b_kv, self.d_model = b_q, b_kv, d_model
+        ### Create learnable vector embeddings for each feature in the Query  block a
+        self.emb_q = self.dist.normal(0,1,shape=(self.b_q,self.d_model), sample=sample, name=f'attention_q_{name}', seed = seed)
+
+
+        ### Create learnable vector embeddings for each feature in the Key/ Value block 
+        self.emb_kv= self.dist.normal(0,1,shape=(self.b_kv,d_model),sample=sample, name=f'attention_kv_{name}', seed = seed)    
+
+        ### Define linear layers to project embeddings into Query, Key, and Value spaces
+        self.q_k_v_proj = self.dist.normal(0,1,shape=(self.d_model,self.d_model,3),sample=sample, name=f'attention_q_k_v_{name}', seed = seed)  
+
+        ### Define a final output layer to map the attention context to the desired matrix shape.
+        self.out = self.dist.normal(0,1,shape=(self.d_model * self.b_q, self.b_q * self.b_kv),  sample=sample, name=f'attention_out_{name}', seed = seed) 
+
+        # Performs the attention mechanism
+        ## three dense (linear) layers : Q = X W_Q,  K = X W_K, V = X W_V
+        ## the three projection layers in attention do not have activation functions. They are purely linear transformations.
+        Q = self.layer_linear(X = self.emb_q, dist = self.q_k_v_proj[:,:,0])
+        K = self.layer_linear(X = self.emb_kv, dist = self.q_k_v_proj[:,:,1])
+        V = self.layer_linear(X = self.emb_kv, dist = self.q_k_v_proj[:,:,2])    
+
+        ## Attention mechanism: 
+        ### Calculate dot-product similarity scores between all Queries and Keys. Scale for stability.
+        scores =jnp.matmul(Q, K.T) / jnp.sqrt(self.d_model ** 0.5)
+
+        #### Convert raw scores into attention weights (probabilities) using softmax.
+        attn = jax.nn.softmax(scores, axis=-1)
+
+        ####Compute the context vector as a weighted average of the Value vectors. Now we normalize the similarity scores across all possible `j` for each `i`
+        context = jnp.matmul(attn, V)  
+
+        # Reshape and pass through the final output layer to get the L_ij block.
+        return self.layer_linear(context.reshape(1, -1), self.out).reshape(self.b_q, self.b_kv) 
+
+    def layer_toeplitz(self, block_size = 32, sample = False, name = '', seed = None):
+        print("⚠️This function is still in development. Use it with caution. ⚠️")
+        """
+        Models a diagonal covariance block C_ii with a diagonal structure.
+
+        This is the simplest structure, assuming all variables in the block are
+        uncorrelated with each other. It only learns their individual variances.
+
+        Args:
+            block_size (int): The dimension of this square block.
+        """
+        self.b = block_size
+        # Learnable parameter for the base variance (stored in log space for unconstrained optimization).
+        self.log_sigma = self.dist.log_normal(1,sample=sample, name = f'toeplitz_log_sigma_{name}', seed = seed)
+        # Learnable parameter controlling correlation decay.
+        self.raw_alpha = self.dist.normal(0.5,1,sample=sample, name = f'toeplitz_raw_alpha_{name}', seed = seed)
+        # Learnable parameter for the diagonal adjustment.
+        self.log_diag = self.dist.normal(0,1,sample=sample, name = f'toeplitz_log_diag_{name}', seed = seed)
+
+        ## Constructs the Toeplitz matrix from the learned parameters.
+        sigma, alpha = jax.nn.softplus(self.log_sigma) + 1e-6, jax.nn.sigmoid(self.raw_alpha)
+        idx = jnp.arange(self.b)
+
+        # Create the first row of the Toeplitz matrix: [sigma, sigma*alpha, sigma*alpha^2, ...].
+        toeplitz_row = sigma * (alpha ** idx)
+        indices = jnp.abs(idx[:, None] - idx[None, :])
+        # Build the full Toeplitz matrix by indexing the first row with the lag matrix.
+        toeplitz = toeplitz_row[indices]
+        diag = jax.nn.softplus(self.log_diag) + 1e-5
+        result = toeplitz + diag * jnp.eye(self.b)
+        return result
+
+    def layer_compound_symmetry(self, block_size, sample= True, name='', seed = None):
+        """
+        Build a Compound Symmetry covariance block.
+
+        All variables share a common variance (sigma^2) and correlation (rho).
+        C = sigma^2 * [(1 - rho) I + rho 11^T]
+
+        Returns a (size, size) SPD matrix.
+        """
+        print("⚠️This function is still in development. Use it with caution. ⚠️")
+        # Store the block size.
+        self.b = block_size
+
+        # Learnable parameter for the common variance (log space).
+        self.log_sigma = self.dist.log_normal(1, sample = sample, seed = seed, name = f'compound_symmetry_log_sigma_{name}')
+
+        # Learnable parameter for the common correlation (raw, to be mapped to valid range).
+        self.raw_rho = self.dist.normal(0, 1, sample = sample, seed = seed, name = f'compound_symmetry_raw_rho_{name}')
+
+        # Learnable diagonal offsets.
+        self.log_diag = self.dist.normal(0, 1, shape = (self.b,), sample = sample, seed = seed, name = f'compound_symmetry_log_diag_{name}')
+
+        # Constructs the compound symmetry matrix from the learned parameters.
+        # Convert log_sigma to positive sigma.
+        sigma = jax.nn.softplus(self.log_sigma) + 1e-6
+        # Define the mathematically valid range for rho to ensure the matrix is PD.
+        low = -1.0 / (self.b - 1.0) + 1e-6 if self.b > 1 else 0.0
+        high = 0.999
+
+        # Map the output of sigmoid (0, 1) to the valid range (low, high).
+        rho = low + jax.nn.sigmoid(self.raw_rho) * (high - low)
+        # Create identity and all-ones matrices as building blocks.
+        I = jnp.eye(self.b)
+        ones = jnp.ones((self.b, self.b))
+
+        # Construct the matrix using its mathematical formula.
+        comp_sym = sigma * ((1.0 - rho) * I + rho * ones)
+        # Convert log_diag to positive offsets.
+        diag = jax.nn.softplus(self.log_diag) + 1e-5
+
+        # Add the diagonal offsets.
+        return comp_sym +  diag * jnp.eye(self.b)
+ 
+    def layer_diagonal(self, block_size, sample= True, name='', seed = None):
+        """
+        Build a Diagonal covariance block.
+
+        Variables are assumed independent; only individual variances are learned.
+
+        Returns a (size, size) SPD diagonal matrix.
+        """
+        print("⚠️This function is still in development. Use it with caution. ⚠️")
+        # Store the block size.
+        self.b = block_size
+        # Learnable vector of per-variable variances (stored in log space).
+        self.log_variances = self.dist.normal(0, 1, sample = sample, shape = (block_size,), name = f"log_variances_{name}",seed=seed)
+
+        #Constructs the diagonal matrix from the learned variances
+        # Convert log-variances to positive variances.
+        variances = jax.nn.softplus(self.log_variances) + 1e-6
+        # Create a diagonal matrix from the variances vector.
+        return jnp.diag(variances)
+
+    def scaled_dot_product_attention(self, Q, K, V):
+        """Compute scaled dot-product attention."""
+        d_k = Q.shape[-1]
+        scores = jnp.matmul(Q, K.T) / jnp.sqrt(d_k)
+        attn_weights = jax.nn.softmax(scores, axis=-1)
+        return jnp.matmul(attn_weights, V), attn_weights
+
+    def __call__(self, X):
+        """Forward pass through the Bayesian attention mechanism."""
+        Q = self.layer(X, self.q_proj)
+        K = self.layer(X, self.k_proj)
+        V = self.layer(X, self.v_proj)
+
+        attn_output, attn_weights = self.scaled_dot_product_attention(Q, K, V)
+        return attn_output, attn_weights
+        
+    def cov(self,hidden_dim,N,a, b, sample = False):
+        """
+        Creates a Bayesian Neural Network (BNN) with two layers for covariance estimation.
+        The first layer maps the input to a hidden dimension using a normal distribution,
+        and the second layer outputs two values per N (offsets for a and b).
+        Parameters:
+        - hidden_dim (int): The number of hidden units in the first layer.
+        - N (int): The number of data points, which determines the size of the input and output.
+        - a (jnp.ndarray): The first set of offsets for the covariance matrix.
+        - b (jnp.ndarray): The second set of offsets for the covariance matrix.
+        """
+        print("⚠️This function is still in development. Use it with caution. ⚠️")
+        # First layer weights/biases: note these are treated as latent parameters
+        W1 = self.dist.normal(0, 1, shape=(N, hidden_dim), name='W1', sample=sample)
+
+        # Second layer weights/biases
+        W2 = self.dist.normal(0, 1, shape=(hidden_dim, 2), name='W2', sample=sample)
+
+        # Create one-hot encoding for each N (each row is a one–hot vector)
+        X = jnp.eye(N)
+
+        hidden = jnp.tanh(jnp.dot(X, W1))  # shape: (N, hidden_dim)
+
+        # Second layer: output two values per cafe (offsets for a and b)
+        delta = jnp.dot(hidden, W2)        # shape: (N, 2)    
+
+        return deterministic('rf', jnp.stack([a, b]) + delta) 
+
+    def get_rho(self, posterior):
+        """
+
+        """
+        print("⚠️This function is still in development. Use it with caution. ⚠️")
+        a_b = jnp.mean(posterior, axis=0) 
+        N= a_b.shape[0]
+
+        # 1. Compute sample covariance matrix
+        mean_a_b = jnp.mean(a_b, axis=0)  # Mean of [a_cafe, b_cafe]
+
+        centered_data = a_b - mean_a_b  # Center data by subtracting the mean
+
+        cov_sample = jnp.dot(centered_data.T, centered_data) / (N - 1)  # Covariance matrix
+
+        # 2. Extract sigma (standard deviations) from the diagonal of the covariance matrix
+        sigma = jnp.sqrt(jnp.diagonal(cov_sample))  # Extract standard deviations (sqrt of     variance)
+
+        # 3. Compute Rho (correlation matrix)
+        rho = cov_sample / (sigma[:, None] * sigma[None, :])  # Normalize covariance to obtain correlation matrix
+        return rho
+
+    def make_pd_and_cholesky_jax(self, A):
+        A = 0.5 * (A + A.T)
+        jitter = 1e-6
+        A += jitter * jnp.eye(A.shape[-1])
+        return jnp.linalg.cholesky(A)