tf-extra 1.0.0__py3-none-any.whl

tf_extra/__init__.py

@@ -0,0 +1,19 @@
+"""
+TF-Extra: A powerful extension library for TensorFlow/Keras.
+
+Provides MultiOptimizer, custom layers, losses, and callbacks
+that extend TensorFlow's functionality.
+
+Usage:
+    import tf_extra
+    opt = tf_extra.optimizers.MultiOptimizer(optimizers_and_layers)
+"""
+
+__version__ = "1.0.0"
+__author__ = "Your Name"
+__license__ = "MIT"
+
+from tf_extra.tf_extra import optimizers
+from tf_extra.tf_extra import layers
+from tf_extra.tf_extra import losses
+from tf_extra.tf_extra import callbacks

tf_extra/callbacks/__init__.py

@@ -0,0 +1,13 @@
+"""TF-Extra Custom Callbacks Module."""
+
+from tf_extra.callbacks.custom_callbacks import (
+    CosineAnnealingScheduler,
+    GradientAccumulationCallback,
+    TimeLimitCallback,
+)
+
+__all__ = [
+    "CosineAnnealingScheduler",
+    "GradientAccumulationCallback",
+    "TimeLimitCallback",
+]

tf_extra/callbacks/custom_callbacks.py

@@ -0,0 +1,104 @@
+"""
+Custom Keras Callbacks for TF-Extra.
+"""
+
+import tensorflow as tf
+import math
+import time
+
+
+class CosineAnnealingScheduler(tf.keras.callbacks.Callback):
+    """
+    Cosine Annealing Learning Rate Scheduler with Warm Restarts.
+
+    Implements SGDR: Stochastic Gradient Descent with Warm Restarts.
+    The learning rate follows a cosine curve from eta_max to eta_min,
+    resetting every T_0 epochs (multiplied by T_mult after each restart).
+
+    Args:
+        eta_min: Minimum learning rate. Default 1e-7.
+        eta_max: Maximum learning rate. Default 1e-3.
+        T_0: Number of epochs for the first restart. Default 10.
+        T_mult: Factor to increase T_0 after each restart. Default 2.
+        verbose: Whether to print learning rate changes. Default False.
+    """
+
+    def __init__(self, eta_min=1e-7, eta_max=1e-3, T_0=10, T_mult=2, verbose=False, **kwargs):
+        super().__init__(**kwargs)
+        self.eta_min = eta_min
+        self.eta_max = eta_max
+        self.T_0 = T_0
+        self.T_mult = T_mult
+        self.verbose = verbose
+        self.T_i = T_0
+        self.T_cur = 0
+
+    def on_epoch_begin(self, epoch, logs=None):
+        lr = self.eta_min + 0.5 * (self.eta_max - self.eta_min) * (
+            1 + math.cos(math.pi * self.T_cur / self.T_i)
+        )
+        tf.keras.backend.set_value(self.model.optimizer.learning_rate, lr)
+
+        if self.verbose:
+            print(f"\nEpoch {epoch + 1}: CosineAnnealing LR = {lr:.8f}")
+
+        self.T_cur += 1
+        if self.T_cur >= self.T_i:
+            self.T_cur = 0
+            self.T_i = int(self.T_i * self.T_mult)
+
+
+class GradientAccumulationCallback(tf.keras.callbacks.Callback):
+    """
+    Simulates larger batch sizes by accumulating gradients over multiple steps.
+
+    Note: This callback works best with a custom training loop. For
+    model.fit(), it adjusts the effective learning rate to simulate
+    accumulation.
+
+    Args:
+        accumulation_steps: Number of steps to accumulate before updating.
+    """
+
+    def __init__(self, accumulation_steps=4, **kwargs):
+        super().__init__(**kwargs)
+        self.accumulation_steps = accumulation_steps
+
+    def on_train_begin(self, logs=None):
+        # Scale learning rate by accumulation steps
+        current_lr = float(tf.keras.backend.get_value(self.model.optimizer.learning_rate))
+        self._original_lr = current_lr
+        scaled_lr = current_lr * self.accumulation_steps
+        tf.keras.backend.set_value(self.model.optimizer.learning_rate, scaled_lr)
+        print(f"GradientAccumulation: Scaled LR from {current_lr:.6f} to {scaled_lr:.6f}")
+
+    def on_train_end(self, logs=None):
+        tf.keras.backend.set_value(self.model.optimizer.learning_rate, self._original_lr)
+
+
+class TimeLimitCallback(tf.keras.callbacks.Callback):
+    """
+    Stops training after a specified number of seconds.
+
+    Useful for Kaggle competitions or cloud instances with time limits.
+
+    Args:
+        max_seconds: Maximum training time in seconds.
+        verbose: Whether to print when stopping. Default True.
+    """
+
+    def __init__(self, max_seconds, verbose=True, **kwargs):
+        super().__init__(**kwargs)
+        self.max_seconds = max_seconds
+        self.verbose = verbose
+
+    def on_train_begin(self, logs=None):
+        self._start_time = time.time()
+
+    def on_epoch_end(self, epoch, logs=None):
+        elapsed = time.time() - self._start_time
+        if elapsed >= self.max_seconds:
+            if self.verbose:
+                print(f"\nTimeLimitCallback: Stopping training after {elapsed:.1f}s "
+                      f"(limit: {self.max_seconds}s)")
+            self.model.stop_training = True

tf_extra/layers/__init__.py

@@ -0,0 +1,27 @@
+"""TF-Extra Custom Layers Module."""
+
+from tf_extra.layers.custom_layers import (
+    SpatialDropout1D,
+    ChannelAttention,
+    SpatialAttention,
+    CBAM,
+    SqueezeExcitation,
+    MultiHeadSelfAttention,
+    PositionalEncoding,
+    GeM,
+    AdaptiveConcatPool2D,
+    Mish,
+)
+
+__all__ = [
+    "SpatialDropout1D",
+    "ChannelAttention",
+    "SpatialAttention",
+    "CBAM",
+    "SqueezeExcitation",
+    "MultiHeadSelfAttention",
+    "PositionalEncoding",
+    "GeM",
+    "AdaptiveConcatPool2D",
+    "Mish",
+]

tf_extra/layers/custom_layers.py

@@ -0,0 +1,361 @@
+"""
+Custom Keras Layers for TF-Extra.
+
+Provides production-ready layers that are commonly needed but not
+included in the default tf.keras.layers module.
+"""
+
+import tensorflow as tf
+import numpy as np
+import math
+
+
+class SpatialDropout1D(tf.keras.layers.Layer):
+    """
+    Spatial 1D Dropout: drops entire 1D feature maps instead of individual elements.
+    Useful for NLP embeddings and 1D convolutions.
+
+    Args:
+        rate: Float between 0 and 1. Fraction of feature maps to drop.
+    """
+
+    def __init__(self, rate=0.5, **kwargs):
+        super().__init__(**kwargs)
+        self.rate = rate
+        self.dropout = tf.keras.layers.Dropout(rate)
+
+    def call(self, inputs, training=None):
+        # inputs shape: (batch, timesteps, channels)
+        # We want to drop entire channels, so we reshape
+        input_shape = tf.shape(inputs)
+        noise_shape = (input_shape[0], 1, input_shape[2])
+        return tf.nn.dropout(inputs, rate=self.rate, noise_shape=noise_shape) if training else inputs
+
+    def get_config(self):
+        config = super().get_config()
+        config["rate"] = self.rate
+        return config
+
+
+class ChannelAttention(tf.keras.layers.Layer):
+    """
+    Channel Attention Module (from CBAM).
+
+    Applies both average-pooling and max-pooling, passes each through
+    a shared MLP, then combines them with a sigmoid gate.
+
+    Args:
+        reduction_ratio: Integer, reduction ratio for the bottleneck. Default 16.
+    """
+
+    def __init__(self, reduction_ratio=16, **kwargs):
+        super().__init__(**kwargs)
+        self.reduction_ratio = reduction_ratio
+
+    def build(self, input_shape):
+        channels = input_shape[-1]
+        reduced = max(1, channels // self.reduction_ratio)
+        self.fc1 = tf.keras.layers.Dense(reduced, activation="relu", use_bias=False)
+        self.fc2 = tf.keras.layers.Dense(channels, use_bias=False)
+        super().build(input_shape)
+
+    def call(self, inputs):
+        avg_pool = tf.reduce_mean(inputs, axis=[1, 2], keepdims=True)
+        max_pool = tf.reduce_max(inputs, axis=[1, 2], keepdims=True)
+
+        avg_out = self.fc2(self.fc1(avg_pool))
+        max_out = self.fc2(self.fc1(max_pool))
+
+        attention = tf.nn.sigmoid(avg_out + max_out)
+        return inputs * attention
+
+    def get_config(self):
+        config = super().get_config()
+        config["reduction_ratio"] = self.reduction_ratio
+        return config
+
+
+class SpatialAttention(tf.keras.layers.Layer):
+    """
+    Spatial Attention Module (from CBAM).
+
+    Concatenates average-pooled and max-pooled features along the channel
+    axis, then applies a convolution with sigmoid activation.
+
+    Args:
+        kernel_size: Integer, kernel size for the convolution. Default 7.
+    """
+
+    def __init__(self, kernel_size=7, **kwargs):
+        super().__init__(**kwargs)
+        self.kernel_size = kernel_size
+
+    def build(self, input_shape):
+        self.conv = tf.keras.layers.Conv2D(
+            1,
+            kernel_size=self.kernel_size,
+            padding="same",
+            activation="sigmoid",
+            use_bias=False,
+        )
+        super().build(input_shape)
+
+    def call(self, inputs):
+        avg_pool = tf.reduce_mean(inputs, axis=-1, keepdims=True)
+        max_pool = tf.reduce_max(inputs, axis=-1, keepdims=True)
+        concat = tf.concat([avg_pool, max_pool], axis=-1)
+        attention = self.conv(concat)
+        return inputs * attention
+
+    def get_config(self):
+        config = super().get_config()
+        config["kernel_size"] = self.kernel_size
+        return config
+
+
+class CBAM(tf.keras.layers.Layer):
+    """
+    Convolutional Block Attention Module.
+
+    Sequentially applies Channel Attention then Spatial Attention.
+
+    Args:
+        reduction_ratio: Reduction ratio for channel attention. Default 16.
+        kernel_size: Kernel size for spatial attention conv. Default 7.
+    """
+
+    def __init__(self, reduction_ratio=16, kernel_size=7, **kwargs):
+        super().__init__(**kwargs)
+        self.channel_attention = ChannelAttention(reduction_ratio)
+        self.spatial_attention = SpatialAttention(kernel_size)
+        self.reduction_ratio = reduction_ratio
+        self.kernel_size = kernel_size
+
+    def call(self, inputs):
+        x = self.channel_attention(inputs)
+        x = self.spatial_attention(x)
+        return x
+
+    def get_config(self):
+        config = super().get_config()
+        config["reduction_ratio"] = self.reduction_ratio
+        config["kernel_size"] = self.kernel_size
+        return config
+
+
+class SqueezeExcitation(tf.keras.layers.Layer):
+    """
+    Squeeze-and-Excitation block.
+
+    Recalibrates channel-wise feature responses by explicitly modelling
+    interdependencies between channels.
+
+    Args:
+        reduction_ratio: Reduction ratio for the bottleneck. Default 16.
+    """
+
+    def __init__(self, reduction_ratio=16, **kwargs):
+        super().__init__(**kwargs)
+        self.reduction_ratio = reduction_ratio
+
+    def build(self, input_shape):
+        channels = input_shape[-1]
+        reduced = max(1, channels // self.reduction_ratio)
+        self.fc1 = tf.keras.layers.Dense(reduced, activation="relu", use_bias=False)
+        self.fc2 = tf.keras.layers.Dense(channels, activation="sigmoid", use_bias=False)
+        super().build(input_shape)
+
+    def call(self, inputs):
+        se = tf.reduce_mean(inputs, axis=[1, 2], keepdims=True)
+        se = self.fc1(se)
+        se = self.fc2(se)
+        return inputs * se
+
+    def get_config(self):
+        config = super().get_config()
+        config["reduction_ratio"] = self.reduction_ratio
+        return config
+
+
+class MultiHeadSelfAttention(tf.keras.layers.Layer):
+    """
+    Multi-Head Self-Attention layer.
+
+    A clean implementation of scaled dot-product multi-head attention
+    that can be used in custom transformer architectures.
+
+    Args:
+        embed_dim: Dimensionality of the embedding.
+        num_heads: Number of attention heads.
+        dropout_rate: Dropout rate for attention weights. Default 0.0.
+    """
+
+    def __init__(self, embed_dim, num_heads, dropout_rate=0.0, **kwargs):
+        super().__init__(**kwargs)
+        self.embed_dim = embed_dim
+        self.num_heads = num_heads
+        self.dropout_rate = dropout_rate
+
+        if embed_dim % num_heads != 0:
+            raise ValueError(
+                f"embed_dim ({embed_dim}) must be divisible by "
+                f"num_heads ({num_heads})."
+            )
+        self.head_dim = embed_dim // num_heads
+
+    def build(self, input_shape):
+        self.wq = tf.keras.layers.Dense(self.embed_dim)
+        self.wk = tf.keras.layers.Dense(self.embed_dim)
+        self.wv = tf.keras.layers.Dense(self.embed_dim)
+        self.dense = tf.keras.layers.Dense(self.embed_dim)
+        self.dropout = tf.keras.layers.Dropout(self.dropout_rate)
+        super().build(input_shape)
+
+    def _split_heads(self, x, batch_size):
+        x = tf.reshape(x, (batch_size, -1, self.num_heads, self.head_dim))
+        return tf.transpose(x, perm=[0, 2, 1, 3])
+
+    def call(self, inputs, mask=None, training=None):
+        batch_size = tf.shape(inputs)[0]
+
+        q = self._split_heads(self.wq(inputs), batch_size)
+        k = self._split_heads(self.wk(inputs), batch_size)
+        v = self._split_heads(self.wv(inputs), batch_size)
+
+        scale = tf.math.sqrt(tf.cast(self.head_dim, tf.float32))
+        scores = tf.matmul(q, k, transpose_b=True) / scale
+
+        if mask is not None:
+            scores += mask * -1e9
+
+        weights = tf.nn.softmax(scores, axis=-1)
+        weights = self.dropout(weights, training=training)
+
+        attention = tf.matmul(weights, v)
+        attention = tf.transpose(attention, perm=[0, 2, 1, 3])
+        attention = tf.reshape(attention, (batch_size, -1, self.embed_dim))
+
+        return self.dense(attention)
+
+    def get_config(self):
+        config = super().get_config()
+        config["embed_dim"] = self.embed_dim
+        config["num_heads"] = self.num_heads
+        config["dropout_rate"] = self.dropout_rate
+        return config
+
+
+class PositionalEncoding(tf.keras.layers.Layer):
+    """
+    Sinusoidal Positional Encoding layer.
+
+    Adds positional information to embeddings using sine and cosine
+    functions of different frequencies.
+
+    Args:
+        max_len: Maximum sequence length. Default 5000.
+        embed_dim: Embedding dimensionality. If None, inferred from input.
+    """
+
+    def __init__(self, max_len=5000, embed_dim=None, **kwargs):
+        super().__init__(**kwargs)
+        self.max_len = max_len
+        self.embed_dim = embed_dim
+
+    def build(self, input_shape):
+        d_model = self.embed_dim or int(input_shape[-1])
+        pe = np.zeros((self.max_len, d_model), dtype=np.float32)
+        position = np.arange(0, self.max_len)[:, np.newaxis]
+        div_term = np.exp(np.arange(0, d_model, 2) * -(math.log(10000.0) / d_model))
+        pe[:, 0::2] = np.sin(position * div_term)
+        pe[:, 1::2] = np.cos(position * div_term[: d_model // 2 + d_model % 2][:pe[:, 1::2].shape[1]])
+        self.pe = tf.constant(pe[np.newaxis, :, :])  # (1, max_len, d_model)
+        super().build(input_shape)
+
+    def call(self, inputs):
+        seq_len = tf.shape(inputs)[1]
+        return inputs + self.pe[:, :seq_len, :]
+
+    def get_config(self):
+        config = super().get_config()
+        config["max_len"] = self.max_len
+        config["embed_dim"] = self.embed_dim
+        return config
+
+
+class GeM(tf.keras.layers.Layer):
+    """
+    Generalized Mean Pooling (GeM).
+
+    Computes the generalized mean of each feature map. When p=1 it
+    becomes average pooling; as p→∞ it becomes max pooling.
+
+    Args:
+        p_init: Initial value of the learnable exponent p. Default 3.0.
+        eps: Small constant for numerical stability. Default 1e-6.
+    """
+
+    def __init__(self, p_init=3.0, eps=1e-6, **kwargs):
+        super().__init__(**kwargs)
+        self.p_init = p_init
+        self.eps = eps
+
+    def build(self, input_shape):
+        self.p = self.add_weight(
+            name="p", shape=(1,), initializer=tf.keras.initializers.Constant(self.p_init), trainable=True
+        )
+        super().build(input_shape)
+
+    def call(self, inputs):
+        x = tf.clip_by_value(inputs, clip_value_min=self.eps, clip_value_max=tf.float32.max)
+        x = tf.pow(x, self.p)
+        x = tf.reduce_mean(x, axis=[1, 2], keepdims=False)
+        x = tf.pow(x, 1.0 / self.p)
+        return x
+
+    def get_config(self):
+        config = super().get_config()
+        config["p_init"] = self.p_init
+        config["eps"] = self.eps
+        return config
+
+
+class AdaptiveConcatPool2D(tf.keras.layers.Layer):
+    """
+    Adaptive Concatenated Pooling (from fastai).
+
+    Concatenates the results of adaptive average pooling and
+    adaptive max pooling, giving the model access to both statistics.
+
+    Args:
+        output_size: Target spatial size (H, W). Default (1, 1).
+    """
+
+    def __init__(self, output_size=(1, 1), **kwargs):
+        super().__init__(**kwargs)
+        self.output_size = output_size
+
+    def call(self, inputs):
+        avg = tf.reduce_mean(inputs, axis=[1, 2], keepdims=True)
+        mx = tf.reduce_max(inputs, axis=[1, 2], keepdims=True)
+        return tf.concat([avg, mx], axis=-1)
+
+    def get_config(self):
+        config = super().get_config()
+        config["output_size"] = self.output_size
+        return config
+
+
+class Mish(tf.keras.layers.Layer):
+    """
+    Mish activation function: x * tanh(softplus(x)).
+
+    A self-regularized non-monotonic activation function that often
+    outperforms ReLU and Swish.
+    """
+
+    def call(self, inputs):
+        return inputs * tf.math.tanh(tf.math.softplus(inputs))
+
+    def get_config(self):
+        return super().get_config()

tf_extra/losses/__init__.py

@@ -0,0 +1,19 @@
+"""TF-Extra Custom Losses Module."""
+
+from tf_extra.losses.custom_losses import (
+    FocalLoss,
+    DiceLoss,
+    TverskyLoss,
+    LabelSmoothingCrossEntropy,
+    ContrastiveLoss,
+    TripletLoss,
+)
+
+__all__ = [
+    "FocalLoss",
+    "DiceLoss",
+    "TverskyLoss",
+    "LabelSmoothingCrossEntropy",
+    "ContrastiveLoss",
+    "TripletLoss",
+]

tf_extra/losses/custom_losses.py

@@ -0,0 +1,196 @@
+"""
+Custom Loss Functions for TF-Extra.
+
+Provides commonly used loss functions not included in tf.keras.losses.
+"""
+
+import tensorflow as tf
+
+
+class FocalLoss(tf.keras.losses.Loss):
+    """
+    Focal Loss for addressing class imbalance.
+
+    Down-weights well-classified examples and focuses on hard negatives.
+    FL(p_t) = -alpha_t * (1 - p_t)^gamma * log(p_t)
+
+    Args:
+        alpha: Balancing factor. Default 0.25.
+        gamma: Focusing parameter. Default 2.0.
+        from_logits: Whether predictions are logits. Default False.
+    """
+
+    def __init__(self, alpha=0.25, gamma=2.0, from_logits=False, **kwargs):
+        super().__init__(**kwargs)
+        self.alpha = alpha
+        self.gamma = gamma
+        self.from_logits = from_logits
+
+    def call(self, y_true, y_pred):
+        y_true = tf.cast(y_true, tf.float32)
+        if self.from_logits:
+            y_pred = tf.nn.sigmoid(y_pred)
+
+        y_pred = tf.clip_by_value(y_pred, 1e-7, 1.0 - 1e-7)
+        p_t = y_true * y_pred + (1 - y_true) * (1 - y_pred)
+        alpha_t = y_true * self.alpha + (1 - y_true) * (1 - self.alpha)
+        focal_weight = alpha_t * tf.pow(1.0 - p_t, self.gamma)
+
+        bce = -(y_true * tf.math.log(y_pred) + (1 - y_true) * tf.math.log(1 - y_pred))
+        return tf.reduce_mean(focal_weight * bce)
+
+    def get_config(self):
+        config = super().get_config()
+        config.update({"alpha": self.alpha, "gamma": self.gamma, "from_logits": self.from_logits})
+        return config
+
+
+class DiceLoss(tf.keras.losses.Loss):
+    """
+    Dice Loss for segmentation tasks.
+
+    Measures the overlap between predicted and ground truth masks.
+    DL = 1 - (2 * |X ∩ Y| + smooth) / (|X| + |Y| + smooth)
+
+    Args:
+        smooth: Smoothing constant to avoid division by zero. Default 1.0.
+    """
+
+    def __init__(self, smooth=1.0, **kwargs):
+        super().__init__(**kwargs)
+        self.smooth = smooth
+
+    def call(self, y_true, y_pred):
+        y_true = tf.cast(tf.reshape(y_true, [-1]), tf.float32)
+        y_pred = tf.cast(tf.reshape(y_pred, [-1]), tf.float32)
+        intersection = tf.reduce_sum(y_true * y_pred)
+        return 1.0 - (2.0 * intersection + self.smooth) / (
+            tf.reduce_sum(y_true) + tf.reduce_sum(y_pred) + self.smooth
+        )
+
+    def get_config(self):
+        config = super().get_config()
+        config["smooth"] = self.smooth
+        return config
+
+
+class TverskyLoss(tf.keras.losses.Loss):
+    """
+    Tversky Loss — generalization of Dice Loss.
+
+    Allows controlling the penalty for false positives vs false negatives
+    via alpha and beta.
+
+    Args:
+        alpha: Weight for false positives. Default 0.5.
+        beta: Weight for false negatives. Default 0.5.
+        smooth: Smoothing constant. Default 1.0.
+    """
+
+    def __init__(self, alpha=0.5, beta=0.5, smooth=1.0, **kwargs):
+        super().__init__(**kwargs)
+        self.alpha = alpha
+        self.beta = beta
+        self.smooth = smooth
+
+    def call(self, y_true, y_pred):
+        y_true = tf.cast(tf.reshape(y_true, [-1]), tf.float32)
+        y_pred = tf.cast(tf.reshape(y_pred, [-1]), tf.float32)
+        tp = tf.reduce_sum(y_true * y_pred)
+        fp = tf.reduce_sum((1 - y_true) * y_pred)
+        fn = tf.reduce_sum(y_true * (1 - y_pred))
+        tversky = (tp + self.smooth) / (tp + self.alpha * fp + self.beta * fn + self.smooth)
+        return 1.0 - tversky
+
+    def get_config(self):
+        config = super().get_config()
+        config.update({"alpha": self.alpha, "beta": self.beta, "smooth": self.smooth})
+        return config
+
+
+class LabelSmoothingCrossEntropy(tf.keras.losses.Loss):
+    """
+    Cross Entropy with Label Smoothing.
+
+    Replaces hard 0/1 targets with smoothed targets to prevent
+    overconfidence and improve generalization.
+
+    Args:
+        smoothing: Label smoothing factor. Default 0.1.
+    """
+
+    def __init__(self, smoothing=0.1, **kwargs):
+        super().__init__(**kwargs)
+        self.smoothing = smoothing
+
+    def call(self, y_true, y_pred):
+        num_classes = tf.cast(tf.shape(y_pred)[-1], tf.float32)
+        y_true = tf.cast(y_true, tf.float32)
+        y_true = y_true * (1.0 - self.smoothing) + self.smoothing / num_classes
+        return tf.keras.losses.categorical_crossentropy(y_true, y_pred)
+
+    def get_config(self):
+        config = super().get_config()
+        config["smoothing"] = self.smoothing
+        return config
+
+
+class ContrastiveLoss(tf.keras.losses.Loss):
+    """
+    Contrastive Loss for Siamese Networks.
+
+    L = y * d^2 + (1-y) * max(margin - d, 0)^2
+
+    Args:
+        margin: Margin for dissimilar pairs. Default 1.0.
+    """
+
+    def __init__(self, margin=1.0, **kwargs):
+        super().__init__(**kwargs)
+        self.margin = margin
+
+    def call(self, y_true, y_pred):
+        y_true = tf.cast(y_true, tf.float32)
+        square_pred = tf.square(y_pred)
+        margin_square = tf.square(tf.maximum(self.margin - y_pred, 0))
+        return tf.reduce_mean(y_true * square_pred + (1 - y_true) * margin_square)
+
+    def get_config(self):
+        config = super().get_config()
+        config["margin"] = self.margin
+        return config
+
+
+class TripletLoss(tf.keras.losses.Loss):
+    """
+    Triplet Loss for metric learning.
+
+    L = max(d(anchor, positive) - d(anchor, negative) + margin, 0)
+
+    Expects y_pred to be a concatenation of [anchor, positive, negative]
+    embeddings along the last axis, each of size embed_dim.
+
+    Args:
+        margin: Margin for the triplet loss. Default 1.0.
+    """
+
+    def __init__(self, margin=1.0, **kwargs):
+        super().__init__(**kwargs)
+        self.margin = margin
+
+    def call(self, y_true, y_pred):
+        # y_pred: (batch, 3 * embed_dim)
+        embed_dim = tf.shape(y_pred)[-1] // 3
+        anchor = y_pred[:, :embed_dim]
+        positive = y_pred[:, embed_dim : 2 * embed_dim]
+        negative = y_pred[:, 2 * embed_dim :]
+
+        pos_dist = tf.reduce_sum(tf.square(anchor - positive), axis=-1)
+        neg_dist = tf.reduce_sum(tf.square(anchor - negative), axis=-1)
+        loss = tf.maximum(pos_dist - neg_dist + self.margin, 0.0)
+        return tf.reduce_mean(loss)
+
+    def get_config(self):
+        config = super().get_config()
+        config["margin"] = self.margin
+        return config

tf_extra/optimizers/__init__.py

@@ -0,0 +1,5 @@
+"""TF-Extra Optimizers Module."""
+
+from tf_extra.optimizers.multi_optimizer import MultiOptimizer
+
+__all__ = ["MultiOptimizer"]

tf_extra/optimizers/multi_optimizer.py

@@ -0,0 +1,234 @@
+"""
+MultiOptimizer: Assign different optimizers to different layers of a Keras model.
+
+This is a drop-in replacement for the deprecated tensorflow_addons MultiOptimizer,
+fully compatible with TensorFlow 2.x (including TF 2.16+ / Keras 3).
+
+Example:
+    import tensorflow as tf
+    import tf_extra
+
+    optimizers = [
+        tf.keras.optimizers.Adam(learning_rate=1e-6),
+        tf.keras.optimizers.Adam(learning_rate=1e-5),
+        tf.keras.optimizers.Adam(learning_rate=1e-4),
+    ]
+
+    optimizers_and_layers = [
+        (optimizers[0], model.layers[:100]),
+        (optimizers[1], model.layers[100:200]),
+        (optimizers[2], model.layers[-1]),
+    ]
+
+    opt = tf_extra.optimizers.MultiOptimizer(optimizers_and_layers)
+    model.compile(optimizer=opt, loss="binary_crossentropy")
+"""
+
+import tensorflow as tf
+import numpy as np
+from typing import List, Tuple, Union, Any, Optional, Sequence
+
+
+class MultiOptimizer(tf.keras.optimizers.Optimizer):
+    """
+    An optimizer that applies different optimizers to different sets of layers.
+
+    Each (optimizer, layers) pair assigns that optimizer to the trainable
+    variables of the given layer(s). Variables not covered by any pair
+    will use the fallback_optimizer (default: Adam with lr=1e-3).
+
+    Args:
+        optimizers_and_layers: A list of (optimizer, layers) tuples.
+            - optimizer: a tf.keras.optimizers.Optimizer instance.
+            - layers: a single tf.keras.layers.Layer, or a list of layers.
+        fallback_optimizer: Optional optimizer for any trainable variables
+            not covered by the pairs above. If None, uncovered variables
+            will raise a warning and use Adam(lr=1e-3).
+        name: Name for this optimizer. Default "MultiOptimizer".
+        **kwargs: Additional keyword arguments passed to the base Optimizer.
+
+    Raises:
+        ValueError: If optimizers_and_layers is empty.
+        TypeError: If an element is not a valid (optimizer, layers) tuple.
+    """
+
+    def __init__(
+        self,
+        optimizers_and_layers: List[
+            Tuple[tf.keras.optimizers.Optimizer, Union[tf.keras.layers.Layer, List[tf.keras.layers.Layer]]]
+        ],
+        fallback_optimizer: Optional[tf.keras.optimizers.Optimizer] = None,
+        name: str = "MultiOptimizer",
+        **kwargs,
+    ):
+        super().__init__(name=name, **kwargs)
+
+        if not optimizers_and_layers:
+            raise ValueError(
+                "`optimizers_and_layers` must be a non-empty list of "
+                "(optimizer, layers) tuples."
+            )
+
+        self._optimizers_and_layers = []
+        self._optimizer_list = []
+
+        for idx, pair in enumerate(optimizers_and_layers):
+            if not isinstance(pair, (list, tuple)) or len(pair) != 2:
+                raise TypeError(
+                    f"Element {idx} of optimizers_and_layers must be a "
+                    f"(optimizer, layers) tuple, got {type(pair)}."
+                )
+            optimizer, layer_or_layers = pair
+
+            if not isinstance(optimizer, tf.keras.optimizers.Optimizer):
+                raise TypeError(
+                    f"Element {idx}: first item must be a "
+                    f"tf.keras.optimizers.Optimizer, got {type(optimizer)}."
+                )
+
+            # Normalize layers to a list
+            if isinstance(layer_or_layers, tf.keras.layers.Layer):
+                layers_list = [layer_or_layers]
+            elif isinstance(layer_or_layers, (list, tuple)):
+                layers_list = list(layer_or_layers)
+            else:
+                raise TypeError(
+                    f"Element {idx}: second item must be a Layer or list of "
+                    f"Layers, got {type(layer_or_layers)}."
+                )
+
+            for layer in layers_list:
+                if not isinstance(layer, tf.keras.layers.Layer):
+                    raise TypeError(
+                        f"Expected tf.keras.layers.Layer, got {type(layer)}."
+                    )
+
+            self._optimizers_and_layers.append((optimizer, layers_list))
+            if optimizer not in self._optimizer_list:
+                self._optimizer_list.append(optimizer)
+
+        # Fallback optimizer
+        if fallback_optimizer is not None:
+            self._fallback_optimizer = fallback_optimizer
+        else:
+            self._fallback_optimizer = tf.keras.optimizers.Adam(learning_rate=1e-3)
+
+        if self._fallback_optimizer not in self._optimizer_list:
+            self._optimizer_list.append(self._fallback_optimizer)
+
+        # Build variable-to-optimizer mapping (populated in _build_mapping)
+        self._var_optimizer_map = None
+
+    # ------------------------------------------------------------------
+    # Internal: build a dict   variable.ref() -> optimizer
+    # ------------------------------------------------------------------
+    def _build_mapping(self, var_list):
+        """Build mapping from variable refs to their assigned optimizer."""
+        self._var_optimizer_map = {}
+        assigned_refs = set()
+
+        for optimizer, layers_list in self._optimizers_and_layers:
+            for layer in layers_list:
+                for var in layer.trainable_variables:
+                    ref = var.ref()
+                    if ref in assigned_refs:
+                        import warnings
+                        warnings.warn(
+                            f"Variable '{var.name}' is assigned to multiple "
+                            f"optimizers. The last assignment wins."
+                        )
+                    self._var_optimizer_map[ref] = optimizer
+                    assigned_refs.add(ref)
+
+        # Assign fallback to any unassigned variables
+        for var in var_list:
+            ref = var.ref()
+            if ref not in self._var_optimizer_map:
+                self._var_optimizer_map[ref] = self._fallback_optimizer
+
+    # ------------------------------------------------------------------
+    # Public API expected by Keras
+    # ------------------------------------------------------------------
+    def apply_gradients(self, grads_and_vars, **kwargs):
+        """
+        Group gradients by their assigned optimizer and apply each group.
+        """
+        grads_and_vars = list(grads_and_vars)
+        all_vars = [v for _, v in grads_and_vars]
+
+        if self._var_optimizer_map is None:
+            self._build_mapping(all_vars)
+
+        # Bucket gradients by optimizer
+        optimizer_grads = {}  # optimizer -> list of (grad, var)
+        for grad, var in grads_and_vars:
+            if grad is None:
+                continue
+            opt = self._var_optimizer_map.get(var.ref(), self._fallback_optimizer)
+            optimizer_grads.setdefault(id(opt), (opt, []))
+            optimizer_grads[id(opt)][1].append((grad, var))
+
+        # Apply each optimizer's gradients
+        results = []
+        for _, (opt, gv_list) in optimizer_grads.items():
+            results.append(opt.apply_gradients(gv_list, **kwargs))
+
+        return results
+
+    # ------------------------------------------------------------------
+    # Convenience: allow model.compile(optimizer=multi_opt) to work
+    # ------------------------------------------------------------------
+    def get_config(self):
+        """Return serializable config dict."""
+        config = super().get_config()
+        opt_layer_configs = []
+        for optimizer, layers_list in self._optimizers_and_layers:
+            opt_layer_configs.append(
+                {
+                    "optimizer": tf.keras.optimizers.serialize(optimizer),
+                    "layer_names": [l.name for l in layers_list],
+                }
+            )
+        config["optimizers_and_layers"] = opt_layer_configs
+        config["fallback_optimizer"] = tf.keras.optimizers.serialize(
+            self._fallback_optimizer
+        )
+        return config
+
+    @property
+    def learning_rate(self):
+        """Return the learning rate of the first optimizer (for logging)."""
+        return self._optimizers_and_layers[0][0].learning_rate
+
+    @learning_rate.setter
+    def learning_rate(self, value):
+        """Set learning rate on ALL sub-optimizers."""
+        for opt, _ in self._optimizers_and_layers:
+            opt.learning_rate = value
+        self._fallback_optimizer.learning_rate = value
+
+    @property
+    def sub_optimizers(self):
+        """Access list of (optimizer, layers) pairs."""
+        return list(self._optimizers_and_layers)
+
+    @property
+    def optimizers(self):
+        """Return flat list of unique optimizers."""
+        return list(self._optimizer_list)
+
+    # ------------------------------------------------------------------
+    # Required abstract methods (no-ops; real work delegated to sub-opts)
+    # ------------------------------------------------------------------
+    def _create_slots(self, var_list):
+        pass
+
+    def _resource_apply_dense(self, grad, handle, apply_state=None):
+        pass
+
+    def _resource_apply_sparse(self, grad, handle, indices, apply_state=None):
+        pass
+
+    def get_gradients(self, loss, params):
+        """Compute gradients — delegates to tf.GradientTape compatible path."""
+        return super().get_gradients(loss, params)

tf_extra-1.0.0.dist-info/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2024 Your Name
+
+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.

tf_extra-1.0.0.dist-info/METADATA

@@ -0,0 +1,213 @@
+Metadata-Version: 2.1
+Name: tf-extra
+Version: 1.0.0
+Summary: A powerful extension library for TensorFlow/Keras — MultiOptimizer, custom layers, losses, and callbacks.
+Home-page: https://github.com/yourusername/tf-extra
+Author: Your Name
+Author-email: Your Name <your.email@example.com>
+License: MIT
+Project-URL: Homepage, https://github.com/yourusername/tf-extra
+Project-URL: Repository, https://github.com/yourusername/tf-extra
+Project-URL: Issues, https://github.com/yourusername/tf-extra/issues
+Keywords: tensorflow,keras,deep-learning,multi-optimizer,attention,focal-loss
+Classifier: Development Status :: 5 - Production/Stable
+Classifier: Intended Audience :: Developers
+Classifier: Intended Audience :: Science/Research
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Operating System :: OS Independent
+Classifier: Programming Language :: Python :: 3
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Requires-Python: >=3.8
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: tensorflow>=2.10.0
+Requires-Dist: numpy>=1.21.0
+Provides-Extra: dev
+Requires-Dist: pytest>=7.0; extra == "dev"
+Requires-Dist: pytest-cov>=4.0; extra == "dev"
+Requires-Dist: flake8>=5.0; extra == "dev"
+Requires-Dist: black>=22.0; extra == "dev"
+Requires-Dist: isort>=5.0; extra == "dev"
+Requires-Dist: mypy>=0.990; extra == "dev"
+
+# 🚀 TF-Extra
+
+**A powerful extension library for TensorFlow/Keras** — the spiritual successor to `tensorflow_addons`.
+
+[![PyPI version](https://badge.fury.io/py/tf-extra.svg)](https://pypi.org/project/tf-extra/)
+[![Python 3.8+](https://img.shields.io/badge/python-3.8+-blue.svg)](https://www.python.org/downloads/)
+[![License: MIT](https://img.shields.io/badge/License-MIT-yellow.svg)](https://opensource.org/licenses/MIT)
+[![TensorFlow 2.10+](https://img.shields.io/badge/TensorFlow-2.10+-orange.svg)](https://www.tensorflow.org/)
+
+---
+
+## ✨ Features
+
+| Module | Components |
+|--------|-----------|
+| **🔧 Optimizers** | `MultiOptimizer` — assign different optimizers to different layers |
+| **📐 Layers** | `CBAM`, `SqueezeExcitation`, `MultiHeadSelfAttention`, `GeM`, `Mish`, `PositionalEncoding`, `AdaptiveConcatPool2D` |
+| **📉 Losses** | `FocalLoss`, `DiceLoss`, `TverskyLoss`, `LabelSmoothingCrossEntropy`, `ContrastiveLoss`, `TripletLoss` |
+| **📞 Callbacks** | `CosineAnnealingScheduler`, `GradientAccumulationCallback`, `TimeLimitCallback` |
+
+---
+
+## 📦 Installation
+
+```bash
+pip install tf-extra
+```
+
+Or install from source:
+
+```bash
+git clone https://github.com/yourusername/tf-extra.git
+cd tf-extra
+pip install -e .
+```
+
+---
+
+## 🚀 Quick Start
+
+### MultiOptimizer (Drop-in replacement for `tfa.optimizers.MultiOptimizer`)
+
+```python
+import tensorflow as tf
+import tf_extra
+
+# Load a pre-trained model
+base_model = tf.keras.applications.ResNet50(weights='imagenet', include_top=False)
+x = tf.keras.layers.GlobalAveragePooling2D()(base_model.output)
+x = tf.keras.layers.Dense(1)(x)
+model = tf.keras.Model(base_model.input, x)
+
+# Define different learning rates for different parts of the model
+optimizers = [
+    tf.keras.optimizers.Adam(learning_rate=1e-6),   # Very slow for early layers
+    tf.keras.optimizers.Adam(learning_rate=1e-5),   # Medium for middle layers
+    tf.keras.optimizers.Adam(learning_rate=1e-4),   # Fast for the head
+]
+
+optimizers_and_layers = [
+    (optimizers[0], base_model.layers[:100]),
+    (optimizers[1], base_model.layers[100:]),
+    (optimizers[2], model.layers[-1]),
+]
+
+opt = tf_extra.optimizers.MultiOptimizer(optimizers_and_layers)
+model.compile(
+    optimizer=opt,
+    loss=tf.keras.losses.BinaryCrossentropy(from_logits=True),
+    metrics=[tf.keras.metrics.BinaryAccuracy(threshold=0, name='accuracy')]
+)
+```
+
+### Focal Loss (for imbalanced datasets)
+
+```python
+import tf_extra
+
+model.compile(
+    optimizer='adam',
+    loss=tf_extra.losses.FocalLoss(alpha=0.25, gamma=2.0),
+    metrics=['accuracy']
+)
+```
+
+### CBAM Attention Layer
+
+```python
+import tf_extra
+
+x = tf.keras.layers.Conv2D(64, 3, padding='same')(inputs)
+x = tf_extra.layers.CBAM(reduction_ratio=16)(x)
+```
+
+### Cosine Annealing with Warm Restarts
+
+```python
+import tf_extra
+
+scheduler = tf_extra.callbacks.CosineAnnealingScheduler(
+    eta_min=1e-7,
+    eta_max=1e-3,
+    T_0=10,
+    T_mult=2
+)
+model.fit(X, y, callbacks=[scheduler])
+```
+
+---
+
+## 📖 Full API Reference
+
+### `tf_extra.optimizers.MultiOptimizer`
+
+```python
+MultiOptimizer(
+    optimizers_and_layers,     # List of (optimizer, layers) tuples
+    fallback_optimizer=None,   # Optional: optimizer for unassigned variables
+    name="MultiOptimizer"
+)
+```
+
+### `tf_extra.layers`
+
+| Layer | Description |
+|-------|-------------|
+| `CBAM(reduction_ratio=16, kernel_size=7)` | Convolutional Block Attention Module |
+| `SqueezeExcitation(reduction_ratio=16)` | SE-Net attention block |
+| `MultiHeadSelfAttention(embed_dim, num_heads)` | Transformer-style self-attention |
+| `GeM(p_init=3.0)` | Generalized Mean Pooling |
+| `Mish()` | Mish activation function |
+| `PositionalEncoding(max_len=5000)` | Sinusoidal positional encoding |
+| `AdaptiveConcatPool2D()` | Concatenated average + max pooling |
+| `SpatialDropout1D(rate=0.5)` | Spatial dropout for 1D data |
+
+### `tf_extra.losses`
+
+| Loss | Description |
+|------|-------------|
+| `FocalLoss(alpha=0.25, gamma=2.0)` | Focal loss for class imbalance |
+| `DiceLoss(smooth=1.0)` | Dice loss for segmentation |
+| `TverskyLoss(alpha=0.5, beta=0.5)` | Generalized Dice loss |
+| `LabelSmoothingCrossEntropy(smoothing=0.1)` | Label smoothing |
+| `ContrastiveLoss(margin=1.0)` | For Siamese networks |
+| `TripletLoss(margin=1.0)` | For metric learning |
+
+### `tf_extra.callbacks`
+
+| Callback | Description |
+|----------|-------------|
+| `CosineAnnealingScheduler(eta_min, eta_max, T_0, T_mult)` | SGDR scheduler |
+| `GradientAccumulationCallback(accumulation_steps=4)` | Gradient accumulation |
+| `TimeLimitCallback(max_seconds)` | Stop training after time limit |
+
+---
+
+## 🧪 Testing
+
+```bash
+pip install -e ".[dev]"
+pytest tests/ -v
+```
+
+---
+
+## 🤝 Contributing
+
+Contributions are welcome! Please read the contributing guidelines and submit a PR.
+
+---
+
+## 📄 License
+
+MIT License — see [LICENSE](LICENSE) for details.
+
+---
+
+## 🙏 Acknowledgements
+
+Inspired by `tensorflow_addons` which is no longer maintained.
+Built to provide the same functionality with modern TensorFlow compatibility.

tf_extra-1.0.0.dist-info/RECORD

@@ -0,0 +1,14 @@
+tf_extra/__init__.py,sha256=tl7OZ4sEFDPOZ1-NL7lshGqMCi7dEBHMzDBcic_xLKY,492
+tf_extra/callbacks/__init__.py,sha256=OxhzxY0CnjrUyI49IJBMzfnr1Vql3z8TtcoXAsWMXQc,288
+tf_extra/callbacks/custom_callbacks.py,sha256=AsiGF7nopKC7bAsyc4pnZlSmdhPuMcjtinV2ur6xTbU,3658
+tf_extra/layers/__init__.py,sha256=xSrCBxNg0GA7tso-lW8TngmBilWVR7_F01J8C1qZxnY,511
+tf_extra/layers/custom_layers.py,sha256=PB8ktQT86AjgcXTAlHSMyIZXTI2jI2Hs7DPf8CWBMZk,11731
+tf_extra/losses/__init__.py,sha256=pjMbfeqnqpceOLmEQD6HY2aVuahp3sMyc2zP1w5scrI,343
+tf_extra/losses/custom_losses.py,sha256=fDDaxCwo_NcXK-51NH7mGLInmA7IHtwqvRNwbgpdgeQ,6250
+tf_extra/optimizers/__init__.py,sha256=M1eLL7IFb6SQXZEHpTz2abXdjby8TFus2nDRfPoEe-k,128
+tf_extra/optimizers/multi_optimizer.py,sha256=oI02Stbg067wyBWU0AdLyLtPlQGMfeALcNKOKNigupg,9199
+tf_extra-1.0.0.dist-info/LICENSE,sha256=F-4b93u0OVrVwGXgMwBRq6MlGyUT9zmre1oh5Gft5Ts,1066
+tf_extra-1.0.0.dist-info/METADATA,sha256=AgpScl2oJkMvnMp3AFfE4XozUaY2wmCge2zkfCLRL8c,6625
+tf_extra-1.0.0.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
+tf_extra-1.0.0.dist-info/top_level.txt,sha256=o7AvNlkeqL9f3EDtsjgPLEBsCmm08OsrI-MMb2DBe4A,9
+tf_extra-1.0.0.dist-info/RECORD,,

tf_extra-1.0.0.dist-info/WHEEL

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

tf_extra-1.0.0.dist-info/top_level.txt

@@ -0,0 +1 @@
+tf_extra