nnetsauce

 21class AdaBoostClassifier(Boosting, ClassifierMixin):
 22    """AdaBoost Classification (SAMME) model class derived from class Boosting
 23
 24    Parameters:
 25
 26        obj: object
 27            any object containing a method fit (obj.fit()) and a method predict
 28            (obj.predict())
 29
 30        n_estimators: int
 31            number of boosting iterations
 32
 33        learning_rate: float
 34            learning rate of the boosting procedure
 35
 36        n_hidden_features: int
 37            number of nodes in the hidden layer
 38
 39        reg_lambda: float
 40            regularization parameter for weights
 41
 42        reg_alpha: float
 43            controls compromize between l1 and l2 norm of weights
 44
 45        activation_name: str
 46            activation function: 'relu', 'tanh', 'sigmoid', 'prelu' or 'elu'
 47
 48        a: float
 49            hyperparameter for 'prelu' or 'elu' activation function
 50
 51        nodes_sim: str
 52            type of simulation for the nodes: 'sobol', 'hammersley', 'halton',
 53            'uniform'
 54
 55        bias: boolean
 56            indicates if the hidden layer contains a bias term (True) or not
 57            (False)
 58
 59        dropout: float
 60            regularization parameter; (random) percentage of nodes dropped out
 61            of the training
 62
 63        direct_link: boolean
 64            indicates if the original predictors are included (True) in model's
 65            fitting or not (False)
 66
 67        n_clusters: int
 68            number of clusters for 'kmeans' or 'gmm' clustering (could be 0:
 69                no clustering)
 70
 71        cluster_encode: bool
 72            defines how the variable containing clusters is treated (default is one-hot)
 73            if `False`, then labels are used, without one-hot encoding
 74
 75        type_clust: str
 76            type of clustering method: currently k-means ('kmeans') or Gaussian
 77            Mixture Model ('gmm')
 78
 79        type_scaling: a tuple of 3 strings
 80            scaling methods for inputs, hidden layer, and clustering respectively
 81            (and when relevant).
 82            Currently available: standardization ('std') or MinMax scaling ('minmax')
 83
 84        col_sample: float
 85            percentage of covariates randomly chosen for training
 86
 87        row_sample: float
 88            percentage of rows chosen for training, by stratified bootstrapping
 89
 90        seed: int
 91            reproducibility seed for nodes_sim=='uniform'
 92
 93        verbose: int
 94            0 for no output, 1 for a progress bar (default is 1)
 95
 96        method: str
 97            type of Adaboost method, 'SAMME' (discrete) or 'SAMME.R' (real)
 98
 99        backend: str
100            "cpu" or "gpu" or "tpu"
101
102    Attributes:
103
104        alpha_: list
105            AdaBoost coefficients alpha_m
106
107        base_learners_: dict
108            a dictionary containing the base learners
109
110    Examples:
111
112    See also [https://github.com/Techtonique/nnetsauce/blob/master/examples/adaboost_classification.py](https://github.com/Techtonique/nnetsauce/blob/master/examples/adaboost_classification.py)
113
114    ```python
115    import nnetsauce as ns
116    import numpy as np
117    from sklearn.datasets import load_breast_cancer
118    from sklearn.linear_model import LogisticRegression
119    from sklearn.model_selection import train_test_split
120    from sklearn import metrics
121    from time import time
122
123    breast_cancer = load_breast_cancer()
124    Z = breast_cancer.data
125    t = breast_cancer.target
126    np.random.seed(123)
127    X_train, X_test, y_train, y_test = train_test_split(Z, t, test_size=0.2)
128
129    # SAMME.R
130    clf = LogisticRegression(solver='liblinear', multi_class = 'ovr',
131                            random_state=123)
132    fit_obj = ns.AdaBoostClassifier(clf,
133                                    n_hidden_features=int(11.22338867),
134                                    direct_link=True,
135                                    n_estimators=250, learning_rate=0.01126343,
136                                    col_sample=0.72684326, row_sample=0.86429443,
137                                    dropout=0.63078613, n_clusters=2,
138                                    type_clust="gmm",
139                                    verbose=1, seed = 123,
140                                    method="SAMME.R")
141
142    start = time()
143    fit_obj.fit(X_train, y_train)
144    print(f"Elapsed {time() - start}")
145
146    start = time()
147    print(fit_obj.score(X_test, y_test))
148    print(f"Elapsed {time() - start}")
149
150    preds = fit_obj.predict(X_test)
151
152    print(metrics.classification_report(preds, y_test))
153
154    ```
155
156    """
157
158    # construct the object -----
159    _estimator_type = "classifier"
160
161    def __init__(
162        self,
163        obj,
164        n_estimators=10,
165        learning_rate=0.1,
166        n_hidden_features=1,
167        reg_lambda=0,
168        reg_alpha=0.5,
169        activation_name="relu",
170        a=0.01,
171        nodes_sim="sobol",
172        bias=True,
173        dropout=0,
174        direct_link=False,
175        n_clusters=2,
176        cluster_encode=True,
177        type_clust="kmeans",
178        type_scaling=("std", "std", "std"),
179        col_sample=1,
180        row_sample=1,
181        seed=123,
182        verbose=1,
183        method="SAMME",
184        backend="cpu",
185    ):
186        self.type_fit = "classification"
187        self.verbose = verbose
188        self.method = method
189        self.reg_lambda = reg_lambda
190        self.reg_alpha = reg_alpha
191
192        super().__init__(
193            obj=obj,
194            n_estimators=n_estimators,
195            learning_rate=learning_rate,
196            n_hidden_features=n_hidden_features,
197            activation_name=activation_name,
198            a=a,
199            nodes_sim=nodes_sim,
200            bias=bias,
201            dropout=dropout,
202            direct_link=direct_link,
203            n_clusters=n_clusters,
204            cluster_encode=cluster_encode,
205            type_clust=type_clust,
206            type_scaling=type_scaling,
207            col_sample=col_sample,
208            row_sample=row_sample,
209            seed=seed,
210            backend=backend,
211        )
212
213        self.alpha_ = []
214        self.base_learners_ = dict.fromkeys(range(n_estimators))
215
216    def fit(self, X, y, sample_weight=None, **kwargs):
217        """Fit Boosting model to training data (X, y).
218
219        Parameters:
220
221            X: {array-like}, shape = [n_samples, n_features]
222                Training vectors, where n_samples is the number
223                of samples and n_features is the number of features.
224
225            y: array-like, shape = [n_samples]
226                Target values.
227
228            **kwargs: additional parameters to be passed to
229                    self.cook_training_set or self.obj.fit
230
231        Returns:
232
233             self: object
234        """
235
236        assert mx.is_factor(y), "y must contain only integers"
237
238        assert self.method in (
239            "SAMME",
240            "SAMME.R",
241        ), "`method` must be either 'SAMME' or 'SAMME.R'"
242
243        assert (self.reg_lambda <= 1) & (
244            self.reg_lambda >= 0
245        ), "must have self.reg_lambda <= 1 &  self.reg_lambda >= 0"
246
247        assert (self.reg_alpha <= 1) & (
248            self.reg_alpha >= 0
249        ), "must have self.reg_alpha <= 1 &  self.reg_alpha >= 0"
250
251        # training
252        n, p = X.shape
253        self.n_classes = len(np.unique(y))
254        self.classes_ = np.unique(y)  # for compatibility with sklearn
255        self.n_classes_ = len(self.classes_)  # for compatibility with sklearn
256
257        if sample_weight is None:
258            w_m = np.repeat(1.0 / n, n)
259        else:
260            w_m = np.asarray(sample_weight)
261
262        base_learner = CustomClassifier(
263            self.obj,
264            n_hidden_features=self.n_hidden_features,
265            activation_name=self.activation_name,
266            a=self.a,
267            nodes_sim=self.nodes_sim,
268            bias=self.bias,
269            dropout=self.dropout,
270            direct_link=self.direct_link,
271            n_clusters=self.n_clusters,
272            type_clust=self.type_clust,
273            type_scaling=self.type_scaling,
274            col_sample=self.col_sample,
275            row_sample=self.row_sample,
276            seed=self.seed,
277        )
278
279        if self.verbose == 1:
280            pbar = Progbar(self.n_estimators)
281
282        if self.method == "SAMME":
283            err_m = 1e6
284            err_bound = 1 - 1 / self.n_classes
285            self.alpha_.append(1.0)
286            x_range_n = range(n)
287
288            for m in range(self.n_estimators):
289                preds = base_learner.fit(
290                    X, y, sample_weight=w_m.ravel(), **kwargs
291                ).predict(X)
292
293                self.base_learners_.update({m: deepcopy(base_learner)})
294
295                cond = [y[i] != preds[i] for i in x_range_n]
296
297                err_m = max(
298                    sum([elt[0] * elt[1] for elt in zip(cond, w_m)]),
299                    2.220446049250313e-16,
300                )  # sum(w_m) == 1
301
302                if self.reg_lambda > 0:
303                    err_m += self.reg_lambda * (
304                        (1 - self.reg_alpha) * 0.5 * sum([x**2 for x in w_m])
305                        + self.reg_alpha * sum([abs(x) for x in w_m])
306                    )
307
308                err_m = min(err_m, err_bound)
309
310                alpha_m = self.learning_rate * log(
311                    (self.n_classes - 1) * (1 - err_m) / err_m
312                )
313
314                self.alpha_.append(alpha_m)
315
316                w_m_temp = [exp(alpha_m * cond[i]) for i in x_range_n]
317
318                sum_w_m = sum(w_m_temp)
319
320                w_m = np.asarray([w_m_temp[i] / sum_w_m for i in x_range_n])
321
322                base_learner.set_params(seed=self.seed + (m + 1) * 1000)
323
324                if self.verbose == 1:
325                    pbar.update(m)
326
327            if self.verbose == 1:
328                pbar.update(self.n_estimators)
329
330            self.n_estimators = len(self.base_learners_)
331            self.classes_ = np.unique(y)
332
333            return self
334
335        if self.method == "SAMME.R":
336            Y = mo.one_hot_encode2(y, self.n_classes)
337
338            if sample_weight is None:
339                w_m = np.repeat(1.0 / n, n)  # (N, 1)
340
341            else:
342                w_m = np.asarray(sample_weight)
343
344            for m in range(self.n_estimators):
345                probs = base_learner.fit(
346                    X, y, sample_weight=w_m.ravel(), **kwargs
347                ).predict_proba(X)
348
349                np.clip(
350                    a=probs, a_min=2.220446049250313e-16, a_max=1.0, out=probs
351                )
352
353                self.base_learners_.update({m: deepcopy(base_learner)})
354
355                w_m *= np.exp(
356                    -1.0
357                    * self.learning_rate
358                    * (1.0 - 1.0 / self.n_classes)
359                    * xlogy(Y, probs).sum(axis=1)
360                )
361
362                w_m /= np.sum(w_m)
363
364                base_learner.set_params(seed=self.seed + (m + 1) * 1000)
365
366                if self.verbose == 1:
367                    pbar.update(m)
368
369            if self.verbose == 1:
370                pbar.update(self.n_estimators)
371
372            self.n_estimators = len(self.base_learners_)
373            self.classes_ = np.unique(y)
374
375            return self
376
377    def predict(self, X, **kwargs):
378        """Predict test data X.
379
380        Parameters:
381
382            X: {array-like}, shape = [n_samples, n_features]
383                Training vectors, where n_samples is the number
384                of samples and n_features is the number of features.
385
386            **kwargs: additional parameters to be passed to
387                  self.cook_test_set
388
389        Returns:
390
391            model predictions: {array-like}
392        """
393        return self.predict_proba(X, **kwargs).argmax(axis=1)
394
395    def predict_proba(self, X, **kwargs):
396        """Predict probabilities for test data X.
397
398        Parameters:
399
400            X: {array-like}, shape = [n_samples, n_features]
401                Training vectors, where n_samples is the number
402                of samples and n_features is the number of features.
403
404            **kwargs: additional parameters to be passed to
405                  self.cook_test_set
406
407        Returns:
408
409            probability estimates for test data: {array-like}
410
411        """
412
413        n_iter = len(self.base_learners_)
414
415        if self.method == "SAMME":
416            ensemble_learner = np.zeros((X.shape[0], self.n_classes))
417
418            # if self.verbose == 1:
419            #    pbar = Progbar(n_iter)
420
421            for idx, base_learner in self.base_learners_.items():
422                preds = base_learner.predict(X, **kwargs)
423
424                ensemble_learner += self.alpha_[idx] * mo.one_hot_encode2(
425                    preds, self.n_classes
426                )
427
428                # if self.verbose == 1:
429                #    pbar.update(idx)
430
431            # if self.verbose == 1:
432            #    pbar.update(n_iter)
433
434            expit_ensemble_learner = expit(ensemble_learner)
435
436            sum_ensemble = expit_ensemble_learner.sum(axis=1)
437
438            return expit_ensemble_learner / sum_ensemble[:, None]
439
440        # if self.method == "SAMME.R":
441        ensemble_learner = 0
442
443        # if self.verbose == 1:
444        #    pbar = Progbar(n_iter)
445
446        for idx, base_learner in self.base_learners_.items():
447            probs = base_learner.predict_proba(X, **kwargs)
448
449            np.clip(a=probs, a_min=2.220446049250313e-16, a_max=1.0, out=probs)
450
451            log_preds_proba = np.log(probs)
452
453            ensemble_learner += (
454                log_preds_proba - log_preds_proba.mean(axis=1)[:, None]
455            )
456
457            # if self.verbose == 1:
458            #    pbar.update(idx)
459
460        ensemble_learner *= self.n_classes - 1
461
462        # if self.verbose == 1:
463        #    pbar.update(n_iter)
464
465        expit_ensemble_learner = expit(ensemble_learner)
466
467        sum_ensemble = expit_ensemble_learner.sum(axis=1)
468
469        return expit_ensemble_learner / sum_ensemble[:, None]
470
471    @property
472    def _estimator_type(self):
473        return "classifier"

 10class AttentionMechanism:
 11    """
 12    A comprehensive class implementing various attention mechanisms
 13    for both univariate time series and tabular data using JAX.
 14
 15    Supported attention types:
 16    - Scaled Dot-Product Attention
 17    - Additive (Bahdanau) Attention
 18    - Multi-Head Attention
 19    - Self-Attention
 20    - Temporal Attention (for sequences)
 21    - Feature Attention (for tabular data)
 22    - Cross-Attention
 23    - Context Vector Attention
 24    """
 25
 26    def __init__(
 27        self,
 28        input_dim: int,
 29        hidden_dim: int = 64,
 30        num_heads: int = 4,
 31        dropout: float = 0.1,
 32        seed: int = 42,
 33    ):
 34        """
 35        Args:
 36            input_dim: Dimension of input features
 37            hidden_dim: Hidden dimension for attention computations
 38            num_heads: Number of attention heads for multi-head attention
 39            dropout: Dropout rate
 40            seed: Random seed for parameter initialization
 41        """
 42        self.input_dim = input_dim
 43        self.hidden_dim = hidden_dim
 44        self.num_heads = num_heads
 45        self.dropout = dropout
 46
 47        # Initialize random key
 48        self.rng = random.PRNGKey(seed)
 49
 50        # Initialize parameters
 51        self.params = self._initialize_parameters()
 52
 53        assert (
 54            hidden_dim % num_heads == 0
 55        ), "hidden_dim must be divisible by num_heads"
 56        self.head_dim = hidden_dim // num_heads
 57
 58    def _initialize_parameters(self) -> Dict:
 59        """Initialize all network parameters using JAX"""
 60        keys = random.split(self.rng, 20)
 61
 62        def init_weight(key, shape):
 63            return random.normal(key, shape) * np.sqrt(2.0 / shape[0])
 64
 65        def init_bias(shape):
 66            return jnp.zeros(shape)
 67
 68        params = {
 69            # Scaled Dot-Product Attention
 70            "query_w": init_weight(keys[0], (self.input_dim, self.hidden_dim)),
 71            "query_b": init_bias((self.hidden_dim,)),
 72            "key_w": init_weight(keys[1], (self.input_dim, self.hidden_dim)),
 73            "key_b": init_bias((self.hidden_dim,)),
 74            "value_w": init_weight(keys[2], (self.input_dim, self.hidden_dim)),
 75            "value_b": init_bias((self.hidden_dim,)),
 76            # Additive Attention
 77            "additive_query_w": init_weight(
 78                keys[3], (self.input_dim, self.hidden_dim)
 79            ),
 80            "additive_query_b": init_bias((self.hidden_dim,)),
 81            "additive_key_w": init_weight(
 82                keys[4], (self.input_dim, self.hidden_dim)
 83            ),
 84            "additive_key_b": init_bias((self.hidden_dim,)),
 85            "additive_v_w": init_weight(keys[5], (self.hidden_dim, 1)),
 86            "additive_v_b": init_bias((1,)),
 87            # Multi-Head Attention
 88            "mha_query_w": init_weight(
 89                keys[6], (self.input_dim, self.hidden_dim)
 90            ),
 91            "mha_query_b": init_bias((self.hidden_dim,)),
 92            "mha_key_w": init_weight(
 93                keys[7], (self.input_dim, self.hidden_dim)
 94            ),
 95            "mha_key_b": init_bias((self.hidden_dim,)),
 96            "mha_value_w": init_weight(
 97                keys[8], (self.input_dim, self.hidden_dim)
 98            ),
 99            "mha_value_b": init_bias((self.hidden_dim,)),
100            "mha_output_w": init_weight(
101                keys[9], (self.hidden_dim, self.hidden_dim)
102            ),
103            "mha_output_b": init_bias((self.hidden_dim,)),
104            # Feature Attention
105            "feature_w1": init_weight(
106                keys[10], (self.input_dim, self.hidden_dim)
107            ),
108            "feature_b1": init_bias((self.hidden_dim,)),
109            "feature_w2": init_weight(
110                keys[11], (self.hidden_dim, self.input_dim)
111            ),
112            "feature_b2": init_bias((self.input_dim,)),
113            # Temporal Attention
114            "temporal_query_w": init_weight(
115                keys[12], (self.input_dim, self.hidden_dim)
116            ),
117            "temporal_query_b": init_bias((self.hidden_dim,)),
118            "temporal_key_w": init_weight(
119                keys[13], (self.input_dim, self.hidden_dim)
120            ),
121            "temporal_key_b": init_bias((self.hidden_dim,)),
122            # Context Vector Attention
123            "context_vector": random.normal(keys[14], (1, 1, self.hidden_dim)),
124            "context_query_w": init_weight(
125                keys[15], (self.hidden_dim, self.hidden_dim)
126            ),
127            "context_query_b": init_bias((self.hidden_dim,)),
128            "context_key_w": init_weight(
129                keys[16], (self.input_dim, self.hidden_dim)
130            ),
131            "context_key_b": init_bias((self.hidden_dim,)),
132            "context_value_w": init_weight(
133                keys[17], (self.input_dim, self.hidden_dim)
134            ),
135            "context_value_b": init_bias((self.hidden_dim,)),
136        }
137
138        return params
139
140    @staticmethod
141    @jit
142    def _apply_dropout(
143        x: jnp.ndarray,
144        key: jax.random.PRNGKey,
145        rate: float,
146        training: bool = True,
147    ) -> jnp.ndarray:
148        """Apply dropout"""
149        if training and rate > 0:
150            keep_prob = 1 - rate
151            mask = random.bernoulli(key, keep_prob, x.shape)
152            return jnp.where(mask, x / keep_prob, 0)
153        return x
154
155    @partial(jit, static_argnums=(0,))
156    def scaled_dot_product_attention(
157        self,
158        query: jnp.ndarray,
159        key: jnp.ndarray,
160        value: jnp.ndarray,
161        params: Dict,
162        mask: Optional[jnp.ndarray] = None,
163        training: bool = False,
164    ) -> Tuple[jnp.ndarray, jnp.ndarray]:
165        """
166        Scaled Dot-Product Attention
167
168        Args:
169            query: (batch_size, seq_len, input_dim) or (batch_size, input_dim)
170            key: (batch_size, seq_len, input_dim)
171            value: (batch_size, seq_len, input_dim)
172            params: Parameter dictionary
173            mask: Optional mask (batch_size, seq_len)
174            training: Whether in training mode
175
176        Returns:
177            context: Attended context vector
178            attention_weights: Attention weights
179        """
180        # Project inputs
181        Q = jnp.dot(query, params["query_w"]) + params["query_b"]
182        K = jnp.dot(key, params["key_w"]) + params["key_b"]
183        V = jnp.dot(value, params["value_w"]) + params["value_b"]
184
185        # Compute attention scores
186        scores = jnp.matmul(Q, jnp.swapaxes(K, -2, -1))
187        scores = scores / jnp.sqrt(self.hidden_dim)
188
189        # Apply mask if provided
190        if mask is not None:
191            scores = jnp.where(mask == 0, -1e9, scores)
192
193        # Compute attention weights
194        attention_weights = jax.nn.softmax(scores, axis=-1)
195
196        # Apply attention to values
197        context = jnp.matmul(attention_weights, V)
198
199        return context, attention_weights
200
201    @partial(jit, static_argnums=(0,))
202    def additive_attention(
203        self,
204        query: jnp.ndarray,
205        key: jnp.ndarray,
206        value: jnp.ndarray,
207        params: Dict,
208        mask: Optional[jnp.ndarray] = None,
209    ) -> Tuple[jnp.ndarray, jnp.ndarray]:
210        """
211        Additive (Bahdanau) Attention
212
213        Args:
214            query: (batch_size, hidden_dim) or (batch_size, 1, hidden_dim)
215            key: (batch_size, seq_len, hidden_dim)
216            value: (batch_size, seq_len, hidden_dim)
217            params: Parameter dictionary
218            mask: Optional mask
219
220        Returns:
221            context: Attended context vector
222            attention_weights: Attention weights
223        """
224        # Ensure query has seq_len dimension
225        if query.ndim == 2:
226            query = jnp.expand_dims(query, axis=1)
227
228        # Project query and key
229        Q = (
230            jnp.dot(query, params["additive_query_w"])
231            + params["additive_query_b"]
232        )
233        K = jnp.dot(key, params["additive_key_w"]) + params["additive_key_b"]
234
235        # Additive attention: score = v^T tanh(W_q Q + W_k K)
236        combined = jnp.tanh(Q + K)
237        scores = (
238            jnp.dot(combined, params["additive_v_w"]) + params["additive_v_b"]
239        )
240        scores = jnp.squeeze(scores, axis=-1)
241
242        # Apply mask if provided
243        if mask is not None:
244            scores = jnp.where(mask == 0, -1e9, scores)
245
246        # Compute attention weights
247        attention_weights = jax.nn.softmax(scores, axis=-1)
248
249        # Apply attention to values
250        context = jnp.matmul(jnp.expand_dims(attention_weights, axis=1), value)
251        context = jnp.squeeze(context, axis=1)
252
253        return context, attention_weights
254
255    @partial(jit, static_argnums=(0,))
256    def multi_head_attention(
257        self,
258        query: jnp.ndarray,
259        key: jnp.ndarray,
260        value: jnp.ndarray,
261        params: Dict,
262        mask: Optional[jnp.ndarray] = None,
263    ) -> Tuple[jnp.ndarray, jnp.ndarray]:
264        """
265        Multi-Head Attention
266
267        Args:
268            query: (batch_size, seq_len_q, input_dim)
269            key: (batch_size, seq_len_k, input_dim)
270            value: (batch_size, seq_len_v, input_dim)
271            params: Parameter dictionary
272            mask: Optional mask
273
274        Returns:
275            output: Multi-head attention output
276            attention_weights: Attention weights from all heads
277        """
278        batch_size = query.shape[0]
279
280        # Project and reshape for multi-head attention
281        Q = jnp.dot(query, params["mha_query_w"]) + params["mha_query_b"]
282        K = jnp.dot(key, params["mha_key_w"]) + params["mha_key_b"]
283        V = jnp.dot(value, params["mha_value_w"]) + params["mha_value_b"]
284
285        Q = Q.reshape(batch_size, -1, self.num_heads, self.head_dim)
286        K = K.reshape(batch_size, -1, self.num_heads, self.head_dim)
287        V = V.reshape(batch_size, -1, self.num_heads, self.head_dim)
288
289        # Transpose for attention: (batch, num_heads, seq_len, head_dim)
290        Q = jnp.transpose(Q, (0, 2, 1, 3))
291        K = jnp.transpose(K, (0, 2, 1, 3))
292        V = jnp.transpose(V, (0, 2, 1, 3))
293
294        # Compute attention scores
295        scores = jnp.matmul(Q, jnp.swapaxes(K, -2, -1)) / jnp.sqrt(
296            self.head_dim
297        )
298
299        # Apply mask if provided
300        if mask is not None:
301            mask_expanded = jnp.expand_dims(jnp.expand_dims(mask, 1), 2)
302            scores = jnp.where(mask_expanded == 0, -1e9, scores)
303
304        # Attention weights
305        attention_weights = jax.nn.softmax(scores, axis=-1)
306
307        # Apply attention to values
308        context = jnp.matmul(attention_weights, V)
309
310        # Reshape back: (batch, seq_len, hidden_dim)
311        context = jnp.transpose(context, (0, 2, 1, 3))
312        context = context.reshape(batch_size, -1, self.hidden_dim)
313
314        # Final linear projection
315        output = (
316            jnp.dot(context, params["mha_output_w"]) + params["mha_output_b"]
317        )
318
319        return output, attention_weights
320
321    @partial(jit, static_argnums=(0,))
322    def self_attention(
323        self, x: jnp.ndarray, params: Dict, mask: Optional[jnp.ndarray] = None
324    ) -> Tuple[jnp.ndarray, jnp.ndarray]:
325        """Self-Attention mechanism"""
326        return self.scaled_dot_product_attention(x, x, x, params, mask)
327
328    @partial(jit, static_argnums=(0,))
329    def temporal_attention(
330        self, x: jnp.ndarray, params: Dict, mask: Optional[jnp.ndarray] = None
331    ) -> Tuple[jnp.ndarray, jnp.ndarray]:
332        """
333        Temporal Attention for time series data
334
335        Args:
336            x: (batch_size, seq_len, input_dim)
337            params: Parameter dictionary
338            mask: Optional mask
339
340        Returns:
341            context: Temporally attended context
342            attention_weights: Temporal attention weights
343        """
344        # Use last time step as query
345        query = x[:, -1:, :]
346
347        Q = (
348            jnp.dot(query, params["temporal_query_w"])
349            + params["temporal_query_b"]
350        )
351        K = jnp.dot(x, params["temporal_key_w"]) + params["temporal_key_b"]
352
353        # Compute attention scores
354        scores = jnp.matmul(Q, jnp.swapaxes(K, -2, -1)) / jnp.sqrt(
355            self.hidden_dim
356        )
357        scores = jnp.squeeze(scores, axis=1)
358
359        # Apply mask if provided
360        if mask is not None:
361            scores = jnp.where(mask == 0, -1e9, scores)
362
363        # Attention weights
364        attention_weights = jax.nn.softmax(scores, axis=-1)
365
366        # Apply attention
367        context = jnp.matmul(jnp.expand_dims(attention_weights, axis=1), x)
368        context = jnp.squeeze(context, axis=1)
369
370        return context, attention_weights
371
372    @partial(jit, static_argnums=(0,))
373    def feature_attention_tabular(
374        self, x: jnp.ndarray, params: Dict
375    ) -> Tuple[jnp.ndarray, jnp.ndarray]:
376        """
377        Feature Attention for tabular data
378
379        Args:
380            x: (batch_size, num_features)
381            params: Parameter dictionary
382
383        Returns:
384            output: Feature-weighted output
385            attention_weights: Feature importance weights
386        """
387        # Compute feature attention weights
388        hidden = jnp.dot(x, params["feature_w1"]) + params["feature_b1"]
389        hidden = jnp.tanh(hidden)
390        logits = jnp.dot(hidden, params["feature_w2"]) + params["feature_b2"]
391        attention_weights = jax.nn.softmax(logits, axis=-1)
392
393        # Apply attention to features
394        output = x * attention_weights
395
396        return output, attention_weights
397
398    @partial(jit, static_argnums=(0,))
399    def context_vector_attention(
400        self, x: jnp.ndarray, params: Dict, mask: Optional[jnp.ndarray] = None
401    ) -> Tuple[jnp.ndarray, jnp.ndarray]:
402        """
403        Context Vector Attention
404        Uses a learnable global context vector as the query.
405
406        Args:
407            x: (batch_size, seq_len, input_dim)
408            params: Parameter dictionary
409            mask: Optional mask (batch_size, seq_len)
410
411        Returns:
412            context: Global context representation (batch_size, hidden_dim)
413            attention_weights: Attention weights (batch_size, seq_len)
414        """
415        batch_size = x.shape[0]
416
417        # Expand context vector for batch
418        context_vec = jnp.broadcast_to(
419            params["context_vector"], (batch_size, 1, self.hidden_dim)
420        )
421
422        # Project context vector and input
423        Q = (
424            jnp.dot(context_vec, params["context_query_w"])
425            + params["context_query_b"]
426        )
427        K = jnp.dot(x, params["context_key_w"]) + params["context_key_b"]
428        V = jnp.dot(x, params["context_value_w"]) + params["context_value_b"]
429
430        # Compute attention scores
431        scores = jnp.matmul(Q, jnp.swapaxes(K, -2, -1)) / jnp.sqrt(
432            self.hidden_dim
433        )
434        scores = jnp.squeeze(scores, axis=1)
435
436        # Apply mask if provided
437        if mask is not None:
438            scores = jnp.where(mask == 0, -1e9, scores)
439
440        # Compute attention weights
441        attention_weights = jax.nn.softmax(scores, axis=-1)
442
443        # Apply attention to values
444        context = jnp.matmul(jnp.expand_dims(attention_weights, axis=1), V)
445        context = jnp.squeeze(context, axis=1)
446
447        return context, attention_weights
448
449    @partial(jit, static_argnums=(0,))
450    def cross_attention(
451        self,
452        query: jnp.ndarray,
453        key_value: jnp.ndarray,
454        params: Dict,
455        mask: Optional[jnp.ndarray] = None,
456    ) -> Tuple[jnp.ndarray, jnp.ndarray]:
457        """Cross-Attention between two different sequences"""
458        return self.scaled_dot_product_attention(
459            query, key_value, key_value, params, mask
460        )
461
462    def __call__(
463        self,
464        x: jnp.ndarray,
465        attention_type: str = "scaled_dot_product",
466        query: Optional[jnp.ndarray] = None,
467        key_value: Optional[jnp.ndarray] = None,
468        mask: Optional[jnp.ndarray] = None,
469        training: bool = False,
470    ) -> Tuple[jnp.ndarray, jnp.ndarray]:
471        """
472        Forward pass with specified attention mechanism
473
474        Args:
475            x: Input tensor
476            attention_type: Type of attention to use
477            query: Optional query for cross-attention
478            key_value: Optional key-value for cross-attention
479            mask: Optional mask
480            training: Whether in training mode
481
482        Returns:
483            output: Attention output
484            attention_weights: Attention weights
485        """
486        if attention_type == "scaled_dot_product":
487            return self.scaled_dot_product_attention(
488                x, x, x, self.params, mask, training
489            )
490        elif attention_type == "additive":
491            return self.additive_attention(
492                x[:, -1:, :], x, x, self.params, mask
493            )
494        elif attention_type == "multi_head":
495            return self.multi_head_attention(x, x, x, self.params, mask)
496        elif attention_type == "self":
497            return self.self_attention(x, self.params, mask)
498        elif attention_type == "temporal":
499            return self.temporal_attention(x, self.params, mask)
500        elif attention_type == "feature":
501            return self.feature_attention_tabular(x, self.params)
502        elif attention_type == "cross":
503            if query is None or key_value is None:
504                raise ValueError(
505                    "Cross-attention requires both query and key_value"
506                )
507            return self.cross_attention(query, key_value, self.params, mask)
508        elif attention_type == "context_vector":
509            return self.context_vector_attention(x, self.params, mask)
510        else:
511            raise ValueError(f"Unknown attention type: {attention_type}")