jax-sklearn 0.1.0

JAX-accelerated machine learning library with scikit-learn compatibility

XLearn: JAX-Accelerated Machine Learning

XLearn is a drop-in replacement for scikit-learn that provides automatic JAX acceleration for machine learning algorithms while maintaining 100% API compatibility.

---

🎉 Release 0.1.0

First public release of XLearn! This initial version provides:

• ✅ Core JAX acceleration for Linear Models, Clustering, and Decomposition
• ✅ Automatic performance optimization with intelligent fallback
• ✅ 100% scikit-learn API compatibility - truly drop-in replacement
• ✅ Comprehensive test suite with Azure Pipelines CI/CD
• ✅ Production-ready proxy system with robust error handling

---

🚀 Key Features

• 🔄 Drop-in Replacement: Use import xlearn as sklearn - no code changes needed
• ⚡ Automatic Acceleration: JAX acceleration is applied automatically when beneficial
• 🧠 Intelligent Fallback: Automatically falls back to NumPy for small datasets
• 🎯 Performance-Aware: Uses heuristics to decide when JAX provides speedup
• 📊 Significant Speedups: 5.53x faster on large datasets (100K+ samples)
• 🔬 High Precision: Maintains numerical accuracy (1e-14 level differences)

---

📈 Performance Highlights

Problem Size	Algorithm	JAX Speedup	Use Case
100K × 1K	LinearRegression	5.53x	Large-scale regression
50 problems	Batch Processing	5.57x	Multiple datasets
15K × 200	PCA	3.2x	Dimensionality reduction
20K × 150	Ridge	4.1x	Regularized regression

---

🛠 Installation

Prerequisites

# Install JAX (choose CPU or GPU version)
pip install jax jaxlib  # CPU version
# OR
pip install jax[gpu]    # GPU version (CUDA)

Install XLearn

# From source (recommended for now)
git clone https://github.com/chenxingqiang/jax-sklearn.git
cd jax-sklearn
pip install -e .

# From PyPI (coming soon)
pip install xlearn

---

🎯 Quick Start

Basic Usage

# Simply replace sklearn with xlearn!
import xlearn as sklearn
from xlearn.linear_model import LinearRegression
from xlearn.cluster import KMeans
from xlearn.decomposition import PCA

# Everything works exactly the same
model = LinearRegression()
model.fit(X, y)
predictions = model.predict(X_test)

# JAX acceleration is applied automatically for large datasets

Performance Comparison

import numpy as np
import time
import xlearn as sklearn

# Generate large dataset
X = np.random.randn(50000, 200)
y = X @ np.random.randn(200) + 0.1 * np.random.randn(50000)

# XLearn automatically uses JAX for large data
model = sklearn.linear_model.LinearRegression()

start_time = time.time()
model.fit(X, y)
print(f"Training time: {time.time() - start_time:.4f}s")
# Output: Training time: 0.1124s (JAX accelerated)

# Check if JAX was used
print(f"Used JAX acceleration: {getattr(model, 'is_using_jax', False)}")

Manual Configuration

import xlearn._jax as jax_config

# Check JAX status
print(f"JAX available: {jax_config.is_jax_available()}")
print(f"JAX platform: {jax_config.get_jax_platform()}")

# Configure JAX settings
jax_config.set_config(enable_jax=True, jax_platform="gpu")

# Use context manager for temporary settings
with jax_config.config_context(enable_jax=False):
    # This will force NumPy implementation
    model = sklearn.linear_model.LinearRegression()
    model.fit(X, y)

---

🔧 Supported Algorithms

✅ Fully Accelerated

• Linear Models: LinearRegression, Ridge, Lasso, ElasticNet
• Clustering: KMeans
• Decomposition: PCA, TruncatedSVD
• Preprocessing: StandardScaler, MinMaxScaler

🚧 In Development

• Ensemble: RandomForest, GradientBoosting
• SVM: Support Vector Machines
• Neural Networks: MLPClassifier, MLPRegressor
• Gaussian Process: GaussianProcessRegressor

📊 All Other Algorithms

All other scikit-learn algorithms are available with automatic fallback to the original NumPy implementation.

---

🎮 When Does XLearn Use JAX?

XLearn automatically decides when to use JAX based on:

Algorithm-Specific Thresholds

# LinearRegression: Uses JAX when complexity > 1e8
# Equivalent to: 100K samples × 1K features, or 32K × 32K, etc.

# KMeans: Uses JAX when complexity > 1e6  
# Equivalent to: 10K samples × 100 features

# PCA: Uses JAX when complexity > 1e7
# Equivalent to: 32K samples × 300 features

Smart Heuristics

• Large datasets: >10K samples typically benefit from JAX
• High-dimensional: >100 features often see speedups
• Iterative algorithms: Clustering, optimization benefit earlier
• Matrix operations: Linear algebra intensive algorithms

---

📊 Benchmarks

Large-Scale Linear Regression

Dataset: 100,000 samples × 1,000 features
┌─────────────┬──────────────┬──────────────┬─────────────┐
│ Implementation │ Training Time │ Memory Usage │ Accuracy    │
├─────────────┼──────────────┼──────────────┼─────────────┤
│ XLearn (JAX) │    0.060s    │    0.37 GB   │ 1e-14 diff  │
│ Scikit-Learn │    0.331s    │    0.37 GB   │ Reference   │
│ Speedup      │   5.53x      │    Same      │ Equivalent  │
└─────────────┴──────────────┴──────────────┴─────────────┘

Batch Processing (50 Problems)

Task: 50 regression problems (5K samples × 100 features each)
┌─────────────┬──────────────┬──────────────┐
│ Method      │ Total Time   │ Speedup      │
├─────────────┼──────────────┼──────────────┤
│ XLearn      │   0.097s     │   5.57x      │
│ Sequential  │   0.540s     │   1.00x      │
└─────────────┴──────────────┴──────────────┘

---

🔬 Technical Details

Architecture

XLearn uses a 5-layer architecture:

1. User Code Layer: 100% scikit-learn API compatibility
2. Compatibility Layer: Transparent proxy system
3. JAX Acceleration Layer: JIT compilation and vectorization
4. Data Management Layer: Automatic NumPy ↔ JAX conversion
5. Hardware Abstraction: CPU/GPU/TPU support

Key Technologies

• JAX: Just-in-time compilation and automatic differentiation
• Proxy Pattern: Transparent algorithm switching
• Performance Heuristics: Intelligent acceleration decisions
• Automatic Fallback: Robust error handling

---

🚨 Requirements

• Python: 3.10+
• JAX: 0.4.20+ (automatically installs jaxlib)
• NumPy: 1.22.0+
• SciPy: 1.8.0+

Optional Dependencies

• CUDA: For GPU acceleration
• TPU: For TPU acceleration (Google Cloud)

---

🐛 Troubleshooting

JAX Not Found

# Check if JAX is available
import xlearn._jax as jax_config
if not jax_config.is_jax_available():
    print("Install JAX: pip install jax jaxlib")

Force NumPy Implementation

import xlearn._jax as jax_config
jax_config.set_config(enable_jax=False)

Debug Performance Decisions

import xlearn._jax as jax_config
jax_config.set_config(debug_performance=True)  # Shows acceleration decisions

---

🤝 Contributing

We welcome contributions! See CONTRIBUTING.md for guidelines.

Development Setup

git clone https://github.com/chenxingqiang/jax-sklearn.git
cd jax-sklearn
python -m venv xlearn-env
source xlearn-env/bin/activate  # Linux/Mac
pip install -e ".[dev]"

Running Tests

pytest tests/
python -m pytest tests/test_jax_acceleration.py -v

---

📄 License

XLearn is released under the BSD 3-Clause License.

---

🙏 Acknowledgments

• JAX Team: For the amazing JAX library
• Scikit-learn Team: For the foundational ML library
• NumPy/SciPy: For numerical computing infrastructure

---

📞 Support

• Issues: GitHub Issues
• Discussions: GitHub Discussions
• Documentation: Full Documentation

---

🚀 Ready to accelerate your machine learning? Try XLearn today!