Metadata-Version: 2.1
Name: tf-extra
Version: 1.2.2
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
Provides-Extra: dev
License-File: LICENSE

# 🚀 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.
